fix: update backends/jax __init__.py to reflect changes to module

ivy-llc · Jul 24, 2024 · 97b6354 · 97b6354
1 parent cd847fd
commit 97b6354
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ivy/functional/backends/jax/__init__.py b/ivy/functional/backends/jax/__init__.py
@@ -238,7 +238,7 @@ def closest_valid_dtype(type=None, /, as_native=False):
 from . import control_flow_ops
 from .control_flow_ops import *
 from . import module
-from .module import Model
+from .module import Module
 
 
 # sub-backends
@@ -249,7 +249,7 @@ def closest_valid_dtype(type=None, /, as_native=False):
 if importlib.util.find_spec("flax"):
     import flax
 
-    NativeModule = Model
+    NativeModule = Module
 elif importlib.util.find_spec("haiku"):
     import haiku as hk