diff --git a/java/com/facebook/yoga/YogaNative.java b/java/com/facebook/yoga/YogaNative.java
index ae1b9721..bc79d20f 100644
--- a/java/com/facebook/yoga/YogaNative.java
+++ b/java/com/facebook/yoga/YogaNative.java
@@ -134,7 +134,8 @@ public class YogaNative {
   static native void jni_YGNodeStyleSetOverflowJNI(long nativePointer, int overflow);
   static native int jni_YGNodeStyleGetDisplayJNI(long nativePointer);
   static native void jni_YGNodeStyleSetDisplayJNI(long nativePointer, int display);
-  public static native void jni_YGNodeStyleSetFlexJNI(long nativePointer, float flex);
+  static native float jni_YGNodeStyleGetFlexJNI(long nativePointer);
+  static native void jni_YGNodeStyleSetFlexJNI(long nativePointer, float flex);
   static native float jni_YGNodeStyleGetFlexGrowJNI(long nativePointer);
   static native void jni_YGNodeStyleSetFlexGrowJNI(long nativePointer, float flexGrow);
   static native float jni_YGNodeStyleGetFlexShrinkJNI(long nativePointer);
diff --git a/java/com/facebook/yoga/YogaNodeJNIBase.java b/java/com/facebook/yoga/YogaNodeJNIBase.java
index d4a1aeb7..5fbb8728 100644
--- a/java/com/facebook/yoga/YogaNodeJNIBase.java
+++ b/java/com/facebook/yoga/YogaNodeJNIBase.java
@@ -313,7 +313,7 @@ public abstract class YogaNodeJNIBase extends YogaNode implements Cloneable {
   }
 
   public float getFlex() {
-    return YogaNative.jni_YGNodeStyleGetFlex(mNativePointer);
+    return useVanillaJNI ? YogaNative.jni_YGNodeStyleGetFlexJNI(mNativePointer) : YogaNative.jni_YGNodeStyleGetFlex(mNativePointer);
   }
 
   public void setFlex(float flex) {
diff --git a/java/jni/YGJNIVanilla.cpp b/java/jni/YGJNIVanilla.cpp
index 7dcce2f8..c3f60160 100644
--- a/java/jni/YGJNIVanilla.cpp
+++ b/java/jni/YGJNIVanilla.cpp
@@ -34,18 +34,10 @@ YG_NODE_JNI_STYLE_PROP(jint, YGPositionType, PositionType);
 YG_NODE_JNI_STYLE_PROP(jint, YGWrap, FlexWrap);
 YG_NODE_JNI_STYLE_PROP(jint, YGOverflow, Overflow);
 YG_NODE_JNI_STYLE_PROP(jint, YGDisplay, Display);
+YG_NODE_JNI_STYLE_PROP(jfloat, float, Flex);
 YG_NODE_JNI_STYLE_PROP(jfloat, float, FlexGrow);
 YG_NODE_JNI_STYLE_PROP(jfloat, float, FlexShrink);
 
-void jni_YGNodeStyleSetFlexJNI(
-    JNIEnv* env,
-    jobject obj,
-    jlong nativePointer,
-    jfloat value) {
-  YGNodeStyleSetFlex(
-      _jlong2YGNodeRef(nativePointer), static_cast<float>(value));
-}
-
 // Yoga specific properties, not compatible with flexbox specification
 YG_NODE_JNI_STYLE_PROP(jfloat, float, AspectRatio);
 
@@ -72,7 +64,6 @@ void registerNativeMethods(
 }
 
 static JNINativeMethod methods[] = {
-    {"jni_YGNodeStyleSetFlexJNI", "(JF)V", (void*) jni_YGNodeStyleSetFlexJNI},
     {"jni_YGNodeStyleGetDirectionJNI",
      "(J)I",
      (void*) jni_YGNodeStyleGetDirectionJNI},
@@ -133,6 +124,8 @@ static JNINativeMethod methods[] = {
     {"jni_YGNodeStyleSetDisplayJNI",
      "(JI)V",
      (void*) jni_YGNodeStyleSetDisplayJNI},
+    {"jni_YGNodeStyleGetFlexJNI", "(J)F", (void*) jni_YGNodeStyleGetFlexJNI},
+    {"jni_YGNodeStyleSetFlexJNI", "(JF)V", (void*) jni_YGNodeStyleSetFlexJNI},
     {"jni_YGNodeStyleGetFlexGrowJNI",
      "(J)F",
      (void*) jni_YGNodeStyleGetFlexGrowJNI},