diff --git a/java/com/facebook/yoga/YogaConfig.java b/java/com/facebook/yoga/YogaConfig.java
index 6a670407..501a3cca 100644
--- a/java/com/facebook/yoga/YogaConfig.java
+++ b/java/com/facebook/yoga/YogaConfig.java
@@ -37,4 +37,8 @@ public abstract class YogaConfig {
   public abstract YogaLogger getLogger();
 
   abstract long getNativePointer();
+
+  public abstract void setUseVanillaJNI(boolean useVanillaJNI);
+
+  public abstract boolean useVanillaJNI();
 }
diff --git a/java/com/facebook/yoga/YogaConfigJNIBase.java b/java/com/facebook/yoga/YogaConfigJNIBase.java
index c093c5c2..3fbd6646 100644
--- a/java/com/facebook/yoga/YogaConfigJNIBase.java
+++ b/java/com/facebook/yoga/YogaConfigJNIBase.java
@@ -10,6 +10,7 @@ public abstract class YogaConfigJNIBase extends YogaConfig {
 
   long mNativePointer;
   private YogaLogger mLogger;
+  private boolean useVanillaJNI = false;
 
   private YogaConfigJNIBase(long nativePointer) {
     if (nativePointer == 0) {
@@ -70,4 +71,14 @@ public abstract class YogaConfigJNIBase extends YogaConfig {
   long getNativePointer() {
     return mNativePointer;
   }
+
+  @Override
+  public void setUseVanillaJNI(boolean useVanillaJNI) {
+    this.useVanillaJNI = useVanillaJNI;
+  }
+
+  @Override
+  public boolean useVanillaJNI() {
+    return this.useVanillaJNI;
+  }
 }
diff --git a/java/com/facebook/yoga/YogaNodeJNIBase.java b/java/com/facebook/yoga/YogaNodeJNIBase.java
index dcf31bbe..87a22dbb 100644
--- a/java/com/facebook/yoga/YogaNodeJNIBase.java
+++ b/java/com/facebook/yoga/YogaNodeJNIBase.java
@@ -46,6 +46,8 @@ public abstract class YogaNodeJNIBase extends YogaNode implements Cloneable {
 
   private boolean mHasNewLayout = true;
 
+  private boolean useVanillaJNI = false;
+
   private YogaNodeJNIBase(long nativePointer) {
     if (nativePointer == 0) {
       throw new IllegalStateException("Failed to allocate native memory");
@@ -59,6 +61,7 @@ public abstract class YogaNodeJNIBase extends YogaNode implements Cloneable {
 
   YogaNodeJNIBase(YogaConfig config) {
     this(YogaNative.jni_YGNodeNewWithConfig(((YogaConfigJNIBase)config).mNativePointer));
+    this.useVanillaJNI = config.useVanillaJNI();
   }
 
   public void reset() {
@@ -284,7 +287,11 @@ public abstract class YogaNodeJNIBase extends YogaNode implements Cloneable {
   }
 
   public void setFlex(float flex) {
-    YogaNative.jni_YGNodeStyleSetFlex(mNativePointer, flex);
+    if (useVanillaJNI) {
+      YogaNative.jni_YGNodeStyleSetFlexJNI(mNativePointer, flex);
+    } else {
+      YogaNative.jni_YGNodeStyleSetFlex(mNativePointer, flex);
+    }
   }
 
   public float getFlexGrow() {