add dtype pp rule

ASEM000 · Apr 13, 2024 · bfb8031 · bfb8031
1 parent cbd8a18
commit bfb8031
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -39,7 +39,7 @@
   assert is_masked(masked_tree[0]) is True
   ```
 
-- Add `dataclasses` rule for `tree_{repr,str}`
+- Add `dataclasses`, `dtype` rule for `tree_{repr,str}`
 
 ## V0.12
 

diff --git a/docs/index.rst b/docs/index.rst
@@ -31,6 +31,4 @@ Apache2.0 License.
 Indices
 =======
 
-* :ref:`genindex`
-
-
+* :ref:`genindex`
diff --git a/sepes/_src/tree_pprint.py b/sepes/_src/tree_pprint.py
@@ -669,3 +669,10 @@ def _(node: Any) -> str:
             return global_info
         shard_info = tree_repr(ShapeDTypePP(shard_shape, dtype))
         return f"G:{global_info}\nS:{shard_info}"
+
+    @tree_str.def_type(type(jax.numpy.float32))
+    @tree_repr.def_type(type(jax.numpy.float32))
+    def _(node, **spec: Unpack[PPSpec]) -> str:
+        out = str(node.dtype)
+        out = out.replace("float", "f").replace("int", "i").replace("complex", "c")
+        return f"jax.numpy.{out}"