Method test_buffer_protocol_tf

python/tests/test_array.py:1815–1875

(self)

    @unittest.skipIf(not has_tf, "requires TensorFlow")
    def test_buffer_protocol_tf(self):
        dtypes_list = [
            (
                mx.bool_,
                tf.bool,
                np.bool_,
            ),
            (
                mx.uint8,
                tf.uint8,
                np.uint8,
            ),
            (
                mx.uint16,
                tf.uint16,
                np.uint16,
            ),
            (
                mx.uint32,
                tf.uint32,
                np.uint32,
            ),
            (mx.uint64, tf.uint64, np.uint64),
            (mx.int8, tf.int8, np.int8),
            (mx.int16, tf.int16, np.int16),
            (mx.int32, tf.int32, np.int32),
            (mx.int64, tf.int64, np.int64),
            (mx.float16, tf.float16, np.float16),
            (mx.float32, tf.float32, np.float32),
            (
                mx.complex64,
                tf.complex64,
                np.complex64,
            ),
        ]

        for mlx_dtype, tf_dtype, np_dtype in dtypes_list:
            a_np = np.random.uniform(low=0, high=100, size=(3, 4)).astype(np_dtype)
            a_tf = tf.constant(a_np, dtype=tf_dtype)
            a_mx = mx.array(np.array(a_tf))
            for f in [
                lambda x: x,
                lambda x: tf.transpose(x) if isinstance(x, tf.Tensor) else x.T,
            ]:
                mv_mx = memoryview(f(a_mx))
                mv_tf = memoryview(f(a_tf))
                if (mv_mx.c_contiguous and mv_tf.c_contiguous) or (
                    mv_mx.f_contiguous and mv_tf.f_contiguous
                ):
                    self.assertEqual(
                        mv_mx.strides, mv_tf.strides, f"{mlx_dtype}{tf_dtype}"
                    )
                self.assertEqual(mv_mx.shape, mv_tf.shape, f"{mlx_dtype}{tf_dtype}")
                self.assertFalse(mv_mx.readonly)
                back_to_npy = np.array(mv_mx)
                self.assertEqualArray(
                    back_to_npy,
                    f(a_tf),