(self)
| 437 | self.assertEqual(a.dtype, b.dtype) |
| 438 | self.assertEqual(a.shape, b.shape) |
| 439 | np.testing.assert_array_equal(a, b) |
| 440 | |
| 441 | def test_as_numpy_array(self): |
| 442 | # From a number. |
| 443 | self.assertNumpyArrayEqual(np.ones([], dtype=np.int64), as_numpy_array(1)) |
| 444 | # From a numpy array. |
| 445 | self.assertNumpyArrayEqual( |
| 446 | np.ones([2], dtype=np.float32), as_numpy_array(np.ones([2], dtype=np.float32)) |
| 447 | ) |
| 448 | # From a TF tensor. |
| 449 | self.assertNumpyArrayEqual( |
| 450 | np.ones([3], dtype=np.float16), |
| 451 | as_numpy_array(tf.ones([3], dtype=tf.float16)), |
| 452 | ) |
| 453 | # From a nested structure. |
| 454 | jax.tree.map( |
| 455 | self.assertNumpyArrayEqual, |
| 456 | { |
| 457 | "a": np.ones([1], dtype=np.float32), |
| 458 | "b": [np.array([2], dtype=np.int64), {"c": np.array([[4]], dtype=np.int32)}], |
| 459 | }, |
| 460 | as_numpy_array( |
| 461 | { |
| 462 | "a": jnp.ones([1], dtype=jnp.float32), |
| 463 | "b": [np.asarray([2]), {"c": tf.convert_to_tensor([[4]])}], |
| 464 | } |
| 465 | ), |
| 466 | ) |
| 467 |
nothing calls this directly
no test coverage detected