diff --git a/java/jni/YGJNI.cpp b/java/jni/YGJNI.cpp
index 28011661..c0b1a99b 100644
--- a/java/jni/YGJNI.cpp
+++ b/java/jni/YGJNI.cpp
@@ -75,34 +75,6 @@ const short int LAYOUT_BORDER_START_INDEX = 14;
 
 bool useBatchingForLayoutOutputs;
 
-class PtrJNodeMap {
-  using JNodeArray = JArrayClass<JYogaNode::javaobject>;
-  std::map<YGNodeRef, size_t> ptrsToIdxs_;
-  alias_ref<JNodeArray> javaNodes_;
-
-public:
-  PtrJNodeMap() : ptrsToIdxs_{}, javaNodes_{} {}
-  PtrJNodeMap(
-      alias_ref<JArrayLong> nativePointers,
-      alias_ref<JNodeArray> javaNodes)
-      : javaNodes_{javaNodes} {
-    auto pin = nativePointers->pinCritical();
-    auto ptrs = pin.get();
-    for (size_t i = 0, n = pin.size(); i < n; ++i) {
-      ptrsToIdxs_[(YGNodeRef) ptrs[i]] = i;
-    }
-  }
-
-  local_ref<JYogaNode> ref(YGNodeRef node) {
-    auto idx = ptrsToIdxs_.find(node);
-    if (idx == ptrsToIdxs_.end()) {
-      return local_ref<JYogaNode>{};
-    } else {
-      return javaNodes_->getElement(idx->second);
-    }
-  }
-};
-
 namespace {
 
 union YGNodeContext {
diff --git a/java/jni/YGJTypes.h b/java/jni/YGJTypes.h
index c2c5fc51..e93d74db 100644
--- a/java/jni/YGJTypes.h
+++ b/java/jni/YGJTypes.h
@@ -6,6 +6,11 @@
  */
 #include <fb/fbjni.h>
 #include <yoga/YGValue.h>
+#include <yoga/Yoga.h>
+#include <map>
+
+using namespace facebook::jni;
+using namespace std;
 
 struct JYogaNode : public facebook::jni::JavaClass<JYogaNode> {
   static constexpr auto kJavaDescriptor = "Lcom/facebook/yoga/YogaNodeJNIBase;";
@@ -28,3 +33,31 @@ struct JYogaLogger : public facebook::jni::JavaClass<JYogaLogger> {
       facebook::jni::alias_ref<JYogaLogLevel>,
       jstring);
 };
+
+class PtrJNodeMap {
+  using JNodeArray = JArrayClass<JYogaNode::javaobject>;
+  std::map<YGNodeRef, size_t> ptrsToIdxs_;
+  alias_ref<JNodeArray> javaNodes_;
+
+public:
+  PtrJNodeMap() : ptrsToIdxs_{}, javaNodes_{} {}
+  PtrJNodeMap(
+      alias_ref<JArrayLong> nativePointers,
+      alias_ref<JNodeArray> javaNodes)
+      : javaNodes_{javaNodes} {
+    auto pin = nativePointers->pinCritical();
+    auto ptrs = pin.get();
+    for (size_t i = 0, n = pin.size(); i < n; ++i) {
+      ptrsToIdxs_[(YGNodeRef) ptrs[i]] = i;
+    }
+  }
+
+  local_ref<JYogaNode> ref(YGNodeRef node) {
+    auto idx = ptrsToIdxs_.find(node);
+    if (idx == ptrsToIdxs_.end()) {
+      return local_ref<JYogaNode>{};
+    } else {
+      return javaNodes_->getElement(idx->second);
+    }
+  }
+};
diff --git a/yoga/Yoga.cpp b/yoga/Yoga.cpp
index d41a1201..7baab631 100644
--- a/yoga/Yoga.cpp
+++ b/yoga/Yoga.cpp
@@ -4093,7 +4093,7 @@ void YGNodeCalculateLayoutWithContext(
   marker = nullptr;
 
 #ifdef YG_ENABLE_EVENTS
-  Event::publish<Event::LayoutPassEnd>(node);
+  Event::publish<Event::LayoutPassEnd>(node, {layoutContext});
 #endif
 
   // We want to get rid off `useLegacyStretchBehaviour` from YGConfig. But we
diff --git a/yoga/event/event.h b/yoga/event/event.h
index 578d2f3a..1a7c881f 100644
--- a/yoga/event/event.h
+++ b/yoga/event/event.h
@@ -71,5 +71,10 @@ struct Event::TypedData<Event::NodeDeallocation> {
   YGConfig* config;
 };
 
+template <>
+struct Event::TypedData<Event::LayoutPassEnd> {
+  void* layoutContext;
+};
+
 } // namespace yoga
 } // namespace facebook