feat: fix c_, r_ (#96)

* feat: fix c_, r_

Signed-off-by: nstarman <[email protected]>

src/quaxed/numpy/_core.py

@@ -411,6 +411,7 @@
 from typing import Any, Literal, TypeVar
 
 import jax.numpy as jnp
+from jax._src.numpy.index_tricks import CClass, RClass
 from jaxtyping import ArrayLike
 
 from quaxed._types import DType
@@ -456,6 +457,47 @@ def squeeze(a: ArrayLike, axis: int | tuple[int, ...] | None = None) -> ArrayLik
     return jnp.squeeze(a, axis=axis)
 
 
+class QuaxedCClass(CClass):  # type: ignore[misc]
+    """Quaxed version of `jax.numpy.CClass`.
+
+    Examples
+    --------
+    >>> import quaxed.numpy as jnp
+    >>> x = jnp.asarray(jnp.linspace(0, 11, 11), dtype=int)
+    >>> jnp.c_[x, x].T
+    Array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 11],
+           [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 11]], dtype=int32)
+
+    """
+
+    @quaxify
+    def __getitem__(self, key: Any) -> Any:
+        return super().__getitem__(key)
+
+
+c_ = QuaxedCClass()
+
+
+class QuaxedRClass(RClass):  # type: ignore[misc]
+    """Quaxed version of `jax.numpy.RClass`.
+
+    Examples
+    --------
+    >>> import quaxed.numpy as jnp
+    >>> x = jnp.asarray(jnp.linspace(0, 4, 4), dtype=int)
+    >>> jnp.r_[x, x]
+    Array([0, 1, 2, 4, 0, 1, 2, 4], dtype=int32)
+
+    """
+
+    @quaxify
+    def __getitem__(self, key: Any) -> Any:
+        return super().__getitem__(key)
+
+
+r_ = QuaxedRClass()
+
+
 # =============================================================================
 
 
@@ -502,7 +544,6 @@ def __dir__() -> list[str]:
         "pi",
         "printoptions",
         "promote_types",
-        "r_",
         "s_",
         "set_printoptions",
         "signedinteger",

tests/numpy/test_jax.py

@@ -384,7 +384,6 @@ def test_broadcast_to(x1):
     assert jnp.all(qnp.broadcast_to(x1, (3, 3)) == jnp.broadcast_to(x1, (3, 3)))
 
 
-@pytest.mark.xfail(reason="Not implemented.")
 def test_c_():
     """Test `quaxed.numpy.c_`."""
     assert jnp.all(qnp.c_[1:3, 4:6] == jnp.c_[1:3, 4:6])