Converts `x` to numpy ndarray recursively. Args: x: a jnp array, numpy array, TF/PyTorch Tensor, or a nested structure of arrays or Tensors. Returns: A nested structure with the same structure as `x` but with values converted to numpy array. Raises: NotImplemen
(x: Any)
| 610 | |
| 611 | |
| 612 | def as_numpy_array(x: Any): |
| 613 | """Converts `x` to numpy ndarray recursively. |
| 614 | |
| 615 | Args: |
| 616 | x: a jnp array, numpy array, TF/PyTorch Tensor, or a nested structure of arrays or Tensors. |
| 617 | |
| 618 | Returns: |
| 619 | A nested structure with the same structure as `x` but with values converted to numpy array. |
| 620 | |
| 621 | Raises: |
| 622 | NotImplementedError: If conversion for the input type is unsupported. |
| 623 | """ |
| 624 | if isinstance(x, (numbers.Number, Tensor)): |
| 625 | return np.array(x) |
| 626 | if isinstance(x, np.ndarray): |
| 627 | return x |
| 628 | if hasattr(x, "detach"): |
| 629 | x = x.detach() |
| 630 | if hasattr(x, "numpy"): |
| 631 | return x.numpy() |
| 632 | if isinstance(x, (Mapping, Sequence)): |
| 633 | return jax.tree.map(as_numpy_array, x) |
| 634 | raise NotImplementedError(f"{type(x)}: {x}") |
| 635 | |
| 636 | |
| 637 | def with_sharding_constraint( |