Working on a fork from core

src/fjformer/__init__.py

@@ -1,6 +1,4 @@
-from fjformer.checkpoint import (
-    CheckpointManager as CheckpointManager
-)
+from fjformer.checkpoint import CheckpointManager as CheckpointManager
 
 from fjformer.sharding import (
     get_jax_mesh as get_jax_mesh,
@@ -14,19 +12,20 @@
     tree_path_to_string as tree_path_to_string,
     make_shard_and_gather_fns as make_shard_and_gather_fns,
     with_sharding_constraint as with_sharding_constraint,
-    create_mesh as create_mesh
+    create_mesh as create_mesh,
 )
 
 from fjformer.utils import (
     JaxRNG as JaxRNG,
     GenerateRNG as GenerateRNG,
 )
-from fjformer import monitor
+from fjformer import monitor as monitor
+from fjformer import core as core
 from fjformer import pallas_operations as pallas_operations
 from fjformer import optimizers as optimizers
 from fjformer import linen as linen
 
-__version__ = "0.0.69"
+__version__ = "0.0.70"
 
 __all__ = (
     "JaxRNG",
@@ -46,5 +45,6 @@
     "CheckpointManager",
     "pallas_operations",
     "optimizers",
-    "linen"
+    "linen",
+    "core",
 )

src/fjformer/core/__init__.py

@@ -0,0 +1,17 @@
+from fjformer.core.implicit_array import (
+    ImplicitArray as ImplicitArray,
+    use_implicit_args as use_implicit_args,
+    aux_field as aux_field,
+    UninitializedAval as UninitializedAval,
+    default_handler as default_handler,
+    primitive_handler as primitive_handler,
+    ArrayValue as ArrayValue,
+    materialize_nested as materialize_nested,
+    EmptyNode as EmptyNode,
+)
+
+from fjformer.core.utilities import (
+    freeze_keys as freeze_keys,
+)
+from fjformer.core import symbols as symbols
+from fjformer.core import types as types