# Copyright 2022 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Tests for the open source DTensor Python API."""

import os

from unittest import mock

from absl.testing import parameterized

import numpy as np

# pylint: disable=g-direct-tensorflow-import

from tensorflow.dtensor.python import api

from tensorflow.dtensor.python import d_variable

from tensorflow.dtensor.python import layout as layout_lib

from tensorflow.dtensor.python import numpy_util

from tensorflow.dtensor.python.tests import test_util

from tensorflow.dtensor.python.tests import test_util_ops

from tensorflow.python.eager import backprop

from tensorflow.python.eager import context

from tensorflow.python.eager.polymorphic_function import polymorphic_function

from tensorflow.python.framework import config

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import dtypes

from tensorflow.python.framework import errors_impl

from tensorflow.python.framework import ops

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import array_ops_stack

from tensorflow.python.ops import gen_array_ops

from tensorflow.python.ops import gen_bitwise_ops

from tensorflow.python.ops import gen_io_ops

from tensorflow.python.ops import gen_linalg_ops

from tensorflow.python.ops import gen_list_ops

from tensorflow.python.ops import gen_math_ops

from tensorflow.python.ops import gen_nn_ops

from tensorflow.python.ops import gen_resource_variable_ops

from tensorflow.python.ops import gen_spectral_ops

from tensorflow.python.ops import gen_stateless_random_ops

from tensorflow.python.ops import gen_string_ops

from tensorflow.python.ops import math_ops

from tensorflow.python.ops import nn_ops

from tensorflow.python.ops import random_ops

from tensorflow.python.ops import special_math_ops

from tensorflow.python.ops import stateless_random_ops

from tensorflow.python.ops import variables

from tensorflow.python.platform import test as tf_test

from tensorflow.python.util import nest

# pylint: enable=g-direct-tensorflow-import

# Makes a 2-D mesh with dimensions as, X(2) and Y(4).

_MESH_DIM_X = 'x'

_MESH_DIM_Y = 'y'

_MESH_DIMS = [_MESH_DIM_X, _MESH_DIM_Y]

_MATMUL_IMPLEMENTED = (('_unsharded', 0,

0), ('_a_unsharded_b_contracting', 0,

1), ('_a_unsharded_b_non_contracting', 0,

2), ('_a_non_contracting_b_unsharded', 1,

0), ('_a_contracting_b_unsharded', 2, 0),

('_a_contracting_b_contracting', 2,

1), ('_a_non_contracting_b_contracting', 1,

1), ('_a_non_contracting_b_non_contracting', 1, 2),

('_a_contracting_b_non_contracting', 2, 2))

_BATCH_MATMUL_IMPLEMENTED = (('_unsharded', 0,

0), ('_a_unsharded_b_contracting', 0,

2), ('_a_unsharded_b_non_contracting', 0,

3), ('_a_batch_b_batch', 1, 1),

('_a_non_contracting_b_unsharded', 2,

0), ('_a_contracting_b_unsharded', 3,

0), ('_a_contracting_b_contracting', 3, 2),

('_a_non_contracting_b_contracting', 2,

2), ('_a_non_contracting_b_non_contracting', 2,

3), ('_a_contracting_b_non_contracting', 3,

3), ('_a_unsharded_b_batch', 0,

1), ('_a_batch_b_unsharded', 1, 0),

('_a_batch_b_contracting', 1,

2), ('_a_batch_b_non_contracting', 1,

3), ('_a_non_contracting_b_batch', 2,

1), ('_a_contracting_b_batch', 3, 1))

_MATMUL_TRANSPOSE = (('', False, False), ('_b_transpose', False, True),

('_a_transpose', True, False), ('_a_transpose_b_transpose',

True, True))

Layout = layout_lib.Layout

Mesh = layout_lib.Mesh

UNSHARDED = layout_lib.UNSHARDED

def select_tol(op, mesh, default_tol, low_res_tol):

# Lowers the tol for math_ops.pow,

# nn_ops.log_softmax_v2 and gen_math_ops.tanh due to

# resolution on TPU

if (op not in [

math_ops.pow, nn_ops.log_softmax_v2, gen_math_ops.tanh,

gen_math_ops.acosh, gen_math_ops.asinh, gen_math_ops.digamma,

gen_math_ops.igammac, gen_math_ops.lgamma, gen_math_ops.log1p,

math_ops.xlog1py, gen_math_ops.xlogy, gen_math_ops.zeta, gen_math_ops.tan,

gen_math_ops.sin, gen_math_ops.sinh, math_ops.softplus

]):

return default_tol

if 'TPU' in mesh.local_devices()[0]:

return low_res_tol

else:

return default_tol

def order_broadcastable_operands(op, lhs, rhs):

# Swaps operands lhs and rhs. Assumes lhs is the broadcasting tensor. Due to

# ops only with right broadcastable operand like gen_math_ops.truncate_div

if (op in [gen_math_ops.truncate_div, gen_math_ops.truncate_mod]):

return rhs, lhs

return lhs, rhs

class DTensorSPMDTest(test_util.DTensorBaseTest):

def setUp(self):

super(DTensorSPMDTest, self).setUp()

self.skipForDeviceType(['TPU'],

'all tests require 8 TPU cores.',

unless_device_count_equals_to=8)

global_ids = test_util.create_device_ids_array((2, 4))

local_ids = np.ravel(global_ids).tolist()

mesh_dict = dict()

for device in ('CPU', 'GPU', 'TPU'):

mesh_dict[device] = Mesh(

[_MESH_DIM_X, _MESH_DIM_Y],

global_ids,

local_ids,

test_util.create_device_list((2, 4), device),

use_xla_spmd=test_util.get_use_xla_spmd(device),

)

self.mesh = self.configTestMesh(mesh_dict)

# Creates a bunch of common layouts used by tests later.

# - 0-d

self.scalar_replicated_layout = Layout.replicated(self.mesh, rank=0)

# - 1-d

self.replicated_layout_1d = Layout.replicated(self.mesh, rank=1)

self.first_dimension_sharded_layout_1d = Layout.batch_sharded(

self.mesh, _MESH_DIM_X, rank=1)

# - 2-d

self.replicated_layout_2d = Layout.replicated(self.mesh, rank=2)

self.first_dimension_sharded_layout = Layout.batch_sharded(

self.mesh, _MESH_DIM_X, rank=2)

self.last_dimension_sharded_layout = Layout.inner_sharded(

self.mesh, _MESH_DIM_X, rank=2)

self.layouts_2d = [

self.replicated_layout_2d, self.first_dimension_sharded_layout,

self.last_dimension_sharded_layout

]

# - 3-d

self.replicated_layout_3d = Layout.replicated(self.mesh, rank=3)

self.first_dimension_sharded_layout_3d = Layout.batch_sharded(

self.mesh, _MESH_DIM_X, rank=3)

self.middle_dimension_sharded_layout_3d = Layout(

[layout_lib.UNSHARDED, _MESH_DIM_X, layout_lib.UNSHARDED], self.mesh)

self.last_dimension_sharded_layout_3d = Layout.inner_sharded(

self.mesh, _MESH_DIM_X, rank=3)

self.layouts_3d = [

self.replicated_layout_3d, self.first_dimension_sharded_layout_3d,

self.middle_dimension_sharded_layout_3d,

self.last_dimension_sharded_layout_3d

]

self.shardings = {

'batch': Layout.batch_sharded,

'inner': Layout.inner_sharded

}

@parameterized.named_parameters(

('unsharded_unsharded', [layout_lib.UNSHARDED, layout_lib.UNSHARDED]),

('x_unsharded', [_MESH_DIM_X, layout_lib.UNSHARDED]),

('unsharded_x', [layout_lib.UNSHARDED, _MESH_DIM_X]),

('x,y', [_MESH_DIM_X, _MESH_DIM_Y]),

)

@mock.patch.dict(

os.environ, {'DTENSOR_ENABLE_REPLICATED_SPMD_AS_DEFAULT_TF.MOD': '1'}

)

def testDefaultReplicatedSpmd(self, shard_specs):

if test_util.is_gpu_present():

dtype = dtypes.int32

else:

dtype = dtypes.float32

x = stateless_random_ops.stateless_random_uniform(

shape=[4, 8], seed=[0, 1], maxval=7, dtype=dtype

)

y = constant_op.constant(7, dtype=dtype)

expected_result = math_ops.Mod(x=x, y=y)

expected_layout = Layout.replicated(self.mesh, rank=2)

dtensor_result = math_ops.Mod(

x=api.relayout(x, layout=Layout(shard_specs, self.mesh)),

y=api.relayout(y, layout=Layout([], self.mesh)),

)

self.assertDTensorEqual(expected_result, expected_layout, dtensor_result)

@parameterized.product(

shard_type=['replicated', 'batch_sharded'], full_matrices=[True, False])

def testQR(self, shard_type, full_matrices):

np.random.seed(123)

inputs = constant_op.constant(

np.random.normal(0.0, 1.0, 8 * 9 * 10).reshape([8, 9, 10]),

dtype=dtypes.float32)

expected_result = gen_linalg_ops.qr(

input=inputs, full_matrices=True, name=None)

if shard_type == 'replicated':

layout = self.first_dimension_sharded_layout_3d

else:

layout = self.replicated_layout_3d

inputs = api.relayout(inputs, layout)

got = gen_linalg_ops.qr(

input=inputs, full_matrices=full_matrices, name=None)

self.assertDTensorEqual(expected_result[0], layout, got[0])

self.assertDTensorEqual(expected_result[1], layout, got[1])

def testReduceScatter(self,):

# Generates an AllReduce due to sharding of inner dimensions of Matmul

# and a Scatter due to the Relayout. The AllReduce+Scatter can be combined

# to a single ReduceScatter.

a, b, c = 128, 128, 128

seed = [0, 1]

first_dim_sharded = self.first_dimension_sharded_layout

second_dim_sharded = self.last_dimension_sharded_layout

with api.default_mesh(self.mesh):

m1 = numpy_util.stateless_random_uniform(

layout=second_dim_sharded, shape=[a, b], seed=seed

)

m2 = numpy_util.stateless_random_uniform(

layout=first_dim_sharded, shape=[b, c], seed=seed

)

@polymorphic_function.function

def func():

m3 = math_ops.matmul(m1, m2)

return m3

@polymorphic_function.function

def scattered_func():

m3 = math_ops.matmul(m1, m2)

return api.relayout(m3, self.first_dimension_sharded_layout)

dtensor_result = func()

dtensor_scattered_result = scattered_func()

self.assertDTensorEqual(dtensor_result, self.first_dimension_sharded_layout,

dtensor_scattered_result)

def testReduceScatterLastDimSharded(

self,

):

# ReduceScatter on non-0th dimension which requires a transpose.

a, b, c = 128, 128, 128

seed = [0, 1]

first_dim_sharded = self.first_dimension_sharded_layout

second_dim_sharded = self.last_dimension_sharded_layout

@polymorphic_function.function

def uniform(shape, seed, layout):

return api.relayout(

stateless_random_ops.stateless_random_uniform(shape=shape, seed=seed),

layout=layout,

)

with api.default_mesh(self.mesh):

m1 = uniform(layout=second_dim_sharded, shape=[a, b], seed=seed)

m2 = uniform(layout=first_dim_sharded, shape=[b, c], seed=seed)

@polymorphic_function.function

def func():

m3 = math_ops.matmul(m1, m2)

return m3

@polymorphic_function.function

def scattered_func():

m3 = math_ops.matmul(m1, m2)

return api.relayout(m3, self.last_dimension_sharded_layout)

dtensor_result = func()

dtensor_scattered_result = scattered_func()

self.assertDTensorEqual(

dtensor_result,

self.last_dimension_sharded_layout,

dtensor_scattered_result,

)

@parameterized.named_parameters(

(

'xu_ux',

[_MESH_DIM_X, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, _MESH_DIM_X],

),

(

'ux_xu',

[layout_lib.UNSHARDED, _MESH_DIM_X],

[_MESH_DIM_X, layout_lib.UNSHARDED],

),

(

'yu_uy',

[_MESH_DIM_Y, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, _MESH_DIM_Y],

),

(

'uy_yu',

[layout_lib.UNSHARDED, _MESH_DIM_Y],

[_MESH_DIM_Y, layout_lib.UNSHARDED],

),

)

def testAllToAll2D(self, src_spec, tgt_spec):

a = constant_op.constant(

np.arange(

8 * 8,

).reshape((8, 8)),

dtype=dtypes.float32,

)

sharded_a = numpy_util.pack_numpy(a, layout=Layout(src_spec, self.mesh))

@polymorphic_function.function

def func(a):

return api.relayout(a, Layout(tgt_spec, self.mesh))

dtensor_result = func(sharded_a)

self.assertDTensorEqual(a, Layout(tgt_spec, self.mesh), dtensor_result)

@parameterized.named_parameters(

(

'yuu_uuy',

[_MESH_DIM_Y, layout_lib.UNSHARDED, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, layout_lib.UNSHARDED, _MESH_DIM_Y],

),

(

'xuu_uux',

[_MESH_DIM_X, layout_lib.UNSHARDED, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, layout_lib.UNSHARDED, _MESH_DIM_X],

),

(

'uux_xuu',

[layout_lib.UNSHARDED, layout_lib.UNSHARDED, _MESH_DIM_X],

[_MESH_DIM_X, layout_lib.UNSHARDED, layout_lib.UNSHARDED],

),

(

'xuu_uxu',

[_MESH_DIM_X, layout_lib.UNSHARDED, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, _MESH_DIM_X, layout_lib.UNSHARDED],

),

(

'uxu_xuu',

[layout_lib.UNSHARDED, _MESH_DIM_X, layout_lib.UNSHARDED],

[_MESH_DIM_X, layout_lib.UNSHARDED, layout_lib.UNSHARDED],

),

(

'xuy_uxy',

[_MESH_DIM_X, layout_lib.UNSHARDED, _MESH_DIM_Y],

[layout_lib.UNSHARDED, _MESH_DIM_X, _MESH_DIM_Y],

),

(

'uxy_xuy',

[layout_lib.UNSHARDED, _MESH_DIM_X, _MESH_DIM_Y],

[_MESH_DIM_X, layout_lib.UNSHARDED, _MESH_DIM_Y],

),

(

'xyu_uyx',

[_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, _MESH_DIM_Y, _MESH_DIM_X],

),

# Requires additional transpose

(

'uxu_uux',

[layout_lib.UNSHARDED, _MESH_DIM_X, layout_lib.UNSHARDED],

[layout_lib.UNSHARDED, layout_lib.UNSHARDED, _MESH_DIM_X],

),

(

'uux_uxu',

[layout_lib.UNSHARDED, layout_lib.UNSHARDED, _MESH_DIM_X],

[layout_lib.UNSHARDED, _MESH_DIM_X, layout_lib.UNSHARDED],

),

(

'xyu_xuy',

[_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

[_MESH_DIM_X, layout_lib.UNSHARDED, _MESH_DIM_Y],

),

(

'xuy_xyu',

[_MESH_DIM_X, layout_lib.UNSHARDED, _MESH_DIM_Y],

[_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

),

(

'yxu_yux',

[_MESH_DIM_Y, _MESH_DIM_X, layout_lib.UNSHARDED],

[_MESH_DIM_Y, layout_lib.UNSHARDED, _MESH_DIM_X],

),

(

'yux_yxu',

[_MESH_DIM_Y, layout_lib.UNSHARDED, _MESH_DIM_X],

[_MESH_DIM_Y, _MESH_DIM_X, layout_lib.UNSHARDED],

),

)

def testAllToAll3D(self, src_spec, tgt_spec):

a = constant_op.constant(

np.arange(8 * 8 * 8).reshape((8, 8, 8)), dtype=dtypes.float32

)

sharded_a = numpy_util.pack_numpy(a, layout=Layout(src_spec, self.mesh))

@polymorphic_function.function

def func(a):

return api.relayout(a, Layout(tgt_spec, self.mesh))

dtensor_result = func(sharded_a)

self.assertDTensorEqual(a, Layout(tgt_spec, self.mesh), dtensor_result)

def testExpandDimsDifferentInputAndOutputLayouts(self,):

src_numpy = np.random.uniform(size=[10, 10])

src = constant_op.constant(src_numpy, dtype=dtypes.float32)

expected = array_ops.expand_dims_v2(src, axis=-1)

src = api.relayout(src, self.replicated_layout_2d)

@polymorphic_function.function

def expand_dims_fn(src):

expanded = array_ops.expand_dims_v2(src, axis=-1)

return api.relayout(expanded, self.first_dimension_sharded_layout_3d)

dtensor_result = expand_dims_fn(src)

self.assertDTensorEqual(expected, self.first_dimension_sharded_layout_3d,

dtensor_result)

@polymorphic_function.function

def expand_dims_list_axis_fn(src):

expanded = array_ops.expand_dims_v2(src, axis=[-1])

return api.relayout(expanded, self.first_dimension_sharded_layout_3d)

dtensor_result_2 = expand_dims_list_axis_fn(src)

self.assertDTensorEqual(expected, self.first_dimension_sharded_layout_3d,

dtensor_result_2)

def testPackAndUnpackAssertion(self):

layout = Layout.replicated(self.mesh, rank=3)

# Due to Perf concerns, `pack` does not check the compatibility of

# components and layout. Here, we inject a wrong value components.

with api.default_mesh(self.mesh):

b = api.pack(

[constant_op.constant([[[(x + 1) * 1.0]]]) for x in range(8)],

layout=layout)

assert b.shape == [1, 1, 1]

# `to_numpy` assumes all unpacked tensors are compatible with the

# layout. So, it picks any component to use if that dimension is replicated.

# In this case, it picks the final one.

result_dtensor = numpy_util.to_numpy(b)

self.assertAllEqual(constant_op.constant([[[8.]]]), result_dtensor)

# assertDTensorEqual does more aggressive check, which respects the layout.

with self.assertRaisesRegex(AssertionError, 'Mismatched value'):

self.assertDTensorEqual(constant_op.constant([[[8.]]]), layout, b)

@parameterized.named_parameters(test_util_ops.UNARY_OPS)

def testUnaryOpsWithTwoShardedAndOneReplicatedDimension(self, op):

a = constant_op.constant([[[1.], [2.], [3.], [4.]], [[5.], [6.], [7.],

[8.]]])

assert a.shape == [2, 4, 1]

expected_result = op(a)

layout = Layout([_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED], self.mesh)

a = api.relayout(a, layout)

dtensor_result = op(a)

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-4)

self.assertDTensorEqual(expected_result, layout, dtensor_result, tol=tol)

@parameterized.named_parameters(test_util_ops.UNARY_OPS)

def testUnaryOpsWithFullyReplicatedInputs(self, op):

a = constant_op.constant([[1., 2.], [3., 4.]])

assert a.shape == [2, 2]

expected_result = op(a)

a = api.copy_to_mesh(a, self.replicated_layout_2d)

dtensor_result = op(a)

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-4)

self.assertDTensorEqual(

expected_result, self.replicated_layout_2d, dtensor_result, tol=tol)

@parameterized.named_parameters(test_util_ops.UNARY_OPS)

def testUnaryOpsWithFullyShardedInputs(self, op):

a = constant_op.constant(

np.arange(16).reshape((2, 4, 2)), dtype=dtypes.float32)

expected_result = op(a)

sharded_layout = Layout([_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

self.mesh)

a = api.relayout(a, sharded_layout)

dtensor_result = op(a)

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-4)

self.assertDTensorEqual(

expected_result, sharded_layout, dtensor_result, tol=tol)

@parameterized.named_parameters(test_util_ops.UNARY_OPS)

def testUnaryOpsWithBatchShardedInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-3)

a = constant_op.constant(np.arange(6).reshape((2, 3)), dtype=dtypes.float32)

expected_result = op(a)

a = api.relayout(a, self.first_dimension_sharded_layout)

dtensor_result = op(a)

self.assertDTensorEqual(

expected_result,

self.first_dimension_sharded_layout,

dtensor_result,

tol=tol)

def testInvertOpsWithFullyShardedInputs(self):

# Invert only support int inputs.

op = lambda x: gen_bitwise_ops.invert(x=x, name='Invert')

a = constant_op.constant(

np.arange(16).reshape((2, 4, 2)), dtype=dtypes.int32)

expected_result = op(a)

sharded_layout = Layout([_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

self.mesh)

a = api.relayout(a, sharded_layout)

dtensor_result = op(a)

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-4)

self.assertDTensorEqual(

expected_result, sharded_layout, dtensor_result, tol=tol)

@parameterized.named_parameters(('replicated', layout_lib.UNSHARDED),

('sharded', _MESH_DIM_X))

def testInvertPermutationOp(self, shard):

self.skipForDeviceType(['GPU', 'TPU'],

'Invert Permutation runs in CPU only.')

op_input = constant_op.constant([3, 4, 0, 2, 1, 5])

expected_result = gen_array_ops.invert_permutation(op_input)

# We should always expected the output to be replicated as the

# expander should relayout both inputs and outputs to replicated.

expected_layout = Layout.replicated(self.mesh, rank=1)

self.assertDTensorEqual(

expected_result,

expected_layout,

gen_array_ops.invert_permutation(

api.relayout(op_input, Layout([shard], self.mesh))

),

)

def testErfcInvOpsWithFullyShardedInputs(self):

# By official doc, math_ops.erfcinv is defined on (0, 2]. In addition,

# math_ops.erfcinv internally calls ndtri internally. So to test the op for

# spmd expanding, we call raw op here.

op = lambda x: gen_math_ops.erfinv(x=x, name='erfinv')

a = constant_op.constant(

np.arange(16).reshape((2, 4, 2)) / 30 + 0.1, dtype=dtypes.float32)

expected_result = op(a)

sharded_layout = Layout([_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

self.mesh)

a = api.relayout(a, sharded_layout)

dtensor_result = op(a)

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-4)

self.assertDTensorEqual(

expected_result, sharded_layout, dtensor_result, tol=tol)

def testPopulationCountWithFullyShardedInputs(self):

# By official doc, gen_bitwise_ops.population_count only supports int

# inputs.

op = lambda x: gen_bitwise_ops.population_count(x=x, name='pc')

a = constant_op.constant(

np.arange(16).reshape((2, 4, 2)), dtype=dtypes.int32)

expected_result = op(a)

sharded_layout = Layout([_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

self.mesh)

a = api.relayout(a, sharded_layout)

dtensor_result = op(a)

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, 1e-4)

self.assertDTensorEqual(

expected_result, sharded_layout, dtensor_result, tol=tol)

def testIgammacOpsWithFullyShardedInputs(self):

# Igammac has super low precision on TPU. So we test it as a separated unit

# tests to avoid lower the tol of other tests.

#

# In addition, according to wiki link below, for s=4, all values are not

# inf/nan.

#

# https://en.wikipedia.org/wiki/Incomplete_gamma_function

tol = 1e-2

op = lambda x: gen_math_ops.igammac(4, x)

a = constant_op.constant(

np.arange(16).reshape((2, 4, 2)), dtype=dtypes.float32)

expected_result = op(a)

sharded_layout = Layout([_MESH_DIM_X, _MESH_DIM_Y, layout_lib.UNSHARDED],

self.mesh)

a = api.relayout(a, sharded_layout)

dtensor_result = op(a)

self.assertDTensorEqual(

expected_result, sharded_layout, dtensor_result, tol=tol)

@parameterized.parameters(('replicated',), ('sharded',))

def testBiasAdd2D(self, shard_type):

value = np.array([[1., 2.], [3., 4.]])

bias = np.array([0.1, 0.2])

expected_result = nn_ops.bias_add(value, bias)

if shard_type == 'replicated':

layout = self.replicated_layout_2d

else:

layout = self.first_dimension_sharded_layout

value = api.relayout(value, layout)

bias = api.relayout(bias, self.replicated_layout_1d)

dtensor_result = nn_ops.bias_add(value, bias)

self.assertDTensorEqual(expected_result, layout, dtensor_result)

@parameterized.product(

shard_type=['replicated', 'batch_sharded'],

data_format=['N...C', 'NC...'])

def testBiasAdd4D(self, shard_type, data_format):

value = np.ones(shape=(6, 2, 4, 2), dtype=np.float32)

bias = np.array([0.1, 0.2], dtype=np.float32)

expected_result = nn_ops.bias_add(value, bias, data_format=data_format)

if shard_type == 'replicated':

layout = Layout.replicated(self.mesh, rank=4)

else:

layout = Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=4)

value = api.relayout(value, layout)

bias = api.relayout(bias, self.replicated_layout_1d)

dtensor_result = nn_ops.bias_add(value, bias, data_format=data_format)

self.assertDTensorEqual(expected_result, layout, dtensor_result)

@parameterized.product(

data_format=['N...C', 'NC...'],

bias_sharding=['x', 'y', layout_lib.UNSHARDED],

c_dim_sharding=['x', layout_lib.UNSHARDED])

def testBiasAddDataFormatTest(self, data_format, bias_sharding,

c_dim_sharding):

if data_format == 'N...C':

c_dim = 3

input_sharding = [

layout_lib.UNSHARDED, layout_lib.UNSHARDED, 'y', c_dim_sharding

]

a = np.ones(shape=(1, 1, 4, 4), dtype=np.float32)

layout = Layout(input_sharding, self.mesh)

else:

c_dim = 1

input_sharding = [

layout_lib.UNSHARDED, c_dim_sharding, 'y', layout_lib.UNSHARDED

]

a = np.ones(shape=(1, 4, 4, 1), dtype=np.float32)

layout = Layout(input_sharding, self.mesh)

bias = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)

expected_result = nn_ops.bias_add(a, bias, data_format=data_format)

expected_result_sharding = input_sharding

if c_dim_sharding == layout_lib.UNSHARDED and bias_sharding != 'y':

expected_result_sharding[c_dim] = bias_sharding

expected_layout = Layout(expected_result_sharding, self.mesh)

a = api.relayout(a, layout)

bias = api.relayout(bias, Layout([bias_sharding], self.mesh))

result = nn_ops.bias_add(a, bias=bias, data_format=data_format)

self.assertDTensorEqual(expected_result, expected_layout, result)

@parameterized.parameters(('replicated',), ('batch_sharded',))

def testBiasAddGrad2D(self, shard_type):

value = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)

expected_result = gen_nn_ops.bias_add_grad(out_backprop=value)

if shard_type == 'replicated':

layout = self.replicated_layout_2d

else:

layout = self.first_dimension_sharded_layout

expected_layout = self.replicated_layout_1d

value = api.relayout(value, layout)

dtensor_result = gen_nn_ops.bias_add_grad(out_backprop=value)

self.assertDTensorEqual(expected_result, expected_layout, dtensor_result)

@parameterized.product(

shard_type=['replicated', 'batch_sharded'], data_format=['NHWC', 'NCHW'])

def testBiasAddGrad4D(self, shard_type, data_format):

value = np.ones(shape=(2, 3, 4, 5), dtype=np.float32)

expected_result = gen_nn_ops.bias_add_grad(

out_backprop=value, data_format=data_format)

if shard_type == 'replicated':

layout = Layout.replicated(self.mesh, rank=4)

else:

layout = Layout.batch_sharded(self.mesh, _MESH_DIM_X, rank=4)

expected_layout = self.replicated_layout_1d

value = api.relayout(value, layout)

dtensor_result = gen_nn_ops.bias_add_grad(

out_backprop=value, data_format=data_format)

self.assertDTensorEqual(expected_result, expected_layout, dtensor_result)

@parameterized.named_parameters(test_util_ops.BINARY_FLOAT_OPS)

def testBinaryOpsWithFullyReplicatedInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

a = constant_op.constant([[1., 2.], [3., 4.]])

b = constant_op.constant([[10., 20.], [30., 40.]])

expected_result = op(a, b)

a = api.copy_to_mesh(a, self.replicated_layout_2d)

b = api.copy_to_mesh(b, self.replicated_layout_2d)

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result, self.replicated_layout_2d, dtensor_result, tol=tol)

@parameterized.named_parameters(test_util_ops.BINARY_FLOAT_OPS)

def testBinaryFloatOpsWithFullyShardedInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

a = constant_op.constant(np.arange(8).reshape((2, 4)), dtype=dtypes.float32)

b = constant_op.constant(

np.arange(8).reshape((2, 4)) + 10.0, dtype=dtypes.float32)

expected_result = op(a, b)

sharded_layout_2d = Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)

a = api.relayout(a, sharded_layout_2d)

b = api.relayout(b, sharded_layout_2d)

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result, sharded_layout_2d, dtensor_result, tol=tol)

@parameterized.named_parameters(test_util_ops.BINARY_BOOL_OPS)

def testBinaryBoolOpsWithFullyShardedInputs(self, op):

a = array_ops.reshape(

constant_op.constant(

[True, False, True, False, True, False, True, False]), [2, 4])

b = array_ops.reshape(

constant_op.constant(

[True, True, True, True, False, False, False, False]), [2, 4])

expected_result = op(a, b)

sharded_layout_2d = Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)

a = api.relayout(a, sharded_layout_2d)

b = api.relayout(b, sharded_layout_2d)

dtensor_result = op(a, b)

self.assertDTensorEqual(expected_result, sharded_layout_2d, dtensor_result)

@parameterized.named_parameters(test_util_ops.BINARY_INT_OPS)

def testBinaryIntOpsWithFullyShardedInputs(self, op):

dtype = dtypes.int64

if test_util.is_gpu_present() and op is gen_math_ops.truncate_mod:

dtype = dtypes.int32

a = constant_op.constant(np.arange(8).reshape((2, 4)), dtype=dtype)

b = constant_op.constant(np.arange(8).reshape((2, 4)) + 1, dtype=dtype)

expected_result = op(a, b)

sharded_layout_2d = Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)

a = api.relayout(a, sharded_layout_2d)

b = api.relayout(b, sharded_layout_2d)

dtensor_result = op(a, b)

self.assertDTensorEqual(expected_result, sharded_layout_2d, dtensor_result)

@parameterized.named_parameters(test_util_ops.BINARY_FLOAT_OPS)

def testBinaryFloatOpsWithBatchShardedInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

a = constant_op.constant(

np.array([[1., 2.], [3., 4.]]), dtype=dtypes.float32)

b = constant_op.constant(

np.array([[10., 20.], [30., 40.]]), dtype=dtypes.float32)

expected_result = op(a, b)

a = api.relayout(a, self.first_dimension_sharded_layout)

b = api.relayout(b, self.first_dimension_sharded_layout)

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result,

self.first_dimension_sharded_layout,

dtensor_result,

tol=tol)

@parameterized.named_parameters(test_util_ops.BINARY_INT_OPS)

def testBinaryIntOpsWithBatchShardedInputs(self, op):

dtype = dtypes.int64

if test_util.is_gpu_present() and op is gen_math_ops.truncate_mod:

dtype = dtypes.int32

a = constant_op.constant(np.array([[1, 2], [3, 4]]), dtype=dtype)

b = constant_op.constant(np.array([[5, 6], [7, 4]]), dtype=dtype)

expected_result = op(a, b)

a = api.relayout(a, self.first_dimension_sharded_layout)

b = api.relayout(b, self.first_dimension_sharded_layout)

dtensor_result = op(a, b)

self.assertDTensorEqual(expected_result,

self.first_dimension_sharded_layout, dtensor_result)

@parameterized.named_parameters(

test_util_ops.BINARY_FLOAT_OPS_WITH_BROADCASTING_SUPPORT

)

def testBinaryFloatOpsWithFullyReplicatedBroadcastableInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

# Currently we only support scalar.

a = constant_op.constant(23.4)

b = constant_op.constant([[10., 20.], [30., 40.]])

expected_result = op(a, b)

a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=a.ndim))

b = api.copy_to_mesh(b, Layout.replicated(self.mesh, rank=b.ndim))

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result, self.replicated_layout_2d, dtensor_result, tol=tol)

@parameterized.named_parameters(

test_util_ops.BINARY_INT_OPS_WITH_BROADCASTING_SUPPORT

)

def testBinaryIntOpsWithFullyReplicatedBroadcastableInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

# Currently we only support scalar.

a = constant_op.constant(3)

b = constant_op.constant([[0, 1], [2, 3]])

a, b = order_broadcastable_operands(op, a, b)

expected_result = op(a, b)

a = api.copy_to_mesh(a, Layout.replicated(self.mesh, rank=a.ndim))

b = api.copy_to_mesh(b, Layout.replicated(self.mesh, rank=b.ndim))

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result, self.replicated_layout_2d, dtensor_result, tol=tol)

@parameterized.named_parameters(

test_util_ops.BINARY_FLOAT_OPS_WITH_BROADCASTING_SUPPORT

)

def testBinaryOpsWithFullyShardedBroadcastableInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

# Currently we only support scalar.

a = constant_op.constant(23.4)

b = constant_op.constant(

10.0 * np.arange(8).reshape((2, 4)), dtype=dtypes.float32)

expected_result = op(a, b)

a = api.copy_to_mesh(a, self.scalar_replicated_layout)

sharded_layout_2d = Layout([_MESH_DIM_X, _MESH_DIM_Y], self.mesh)

b = api.relayout(b, sharded_layout_2d)

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result, sharded_layout_2d, dtensor_result, tol=tol)

@parameterized.named_parameters(

test_util_ops.BINARY_FLOAT_OPS_WITH_BROADCASTING_SUPPORT

)

def testBinaryOpsWithBatchShardedBroadcastableInputs(self, op):

tol = select_tol(op, self.mesh, test_util.DEFAULT_TOL, low_res_tol=1e-2)

# Currently we only support scalar.

a = constant_op.constant(23.4)

b = constant_op.constant(

np.array([[10., 20.], [30., 40.]]), dtype=dtypes.float32)

expected_result = op(a, b)

a = api.copy_to_mesh(a, self.scalar_replicated_layout)

b = api.relayout(b, self.first_dimension_sharded_layout)

dtensor_result = op(a, b)

self.assertDTensorEqual(

expected_result,

self.first_dimension_sharded_layout,

dtensor_result,

tol=tol)

@parameterized.named_parameters(

test_util_ops.expand_test_config(

[

{

'testcase_name': 'Concat',

'op': (lambda v: array_ops.concat(values=v, axis=1)),

},

{

'testcase_name':

'ConcatV1',

'op':

(lambda v: gen_array_ops.concat(concat_dim=1, values=v)),

},

{

'testcase_name': 'ConcatV2',

'op': (lambda v: gen_array_ops.concat_v2(values=v, axis=1)),

},

],

[

{

'shard_type': 'replicated',

},

{

'shard_type': 'sharded',

},

{

'shard_type': 'mixed',

},

],

))

def testConcatOpSPMD(self, op, shard_type):

layout_a = self.replicated_layout_2d

layout_b = self.replicated_layout_2d

layout_output = self.replicated_layout_2d

if shard_type == 'sharded':

layout_a = self.first_dimension_sharded_layout

layout_b = self.first_dimension_sharded_layout

layout_output = self.first_dimension_sharded_layout

elif shard_type == 'mixed':

layout_b = self.first_dimension_sharded_layout

layout_output = self.first_dimension_sharded_layout

a = constant_op.constant([[1., 2.], [3., 4.]])

b = constant_op.constant([[1., 2.], [3., 4.]])

expected_result = op([a, b])

with api.default_mesh(self.mesh):

a = api.relayout(a, layout_a)

b = api.relayout(b, layout_b)

c = op([a, b])

self.assertDTensorEqual(expected_result, layout_output, c)

@parameterized.named_parameters([{

'testcase_name': 'ConcatV1',

'op': (lambda v: gen_array_ops.concat(concat_dim=1, values=v))

}, {

'testcase_name': 'ConcatV2',

'op': (lambda v: gen_array_ops.concat_v2(values=v, axis=1))

}])

def testConcatOpShardedOnConcatDim(self, op):

a = constant_op.constant(

np.arange(16).reshape((2, 2, 4)), dtype=dtypes.float32)

b = constant_op.constant(

np.arange(16).reshape((2, 2, 4)), dtype=dtypes.float32)

expected_result = op([a, b])

a_layout = Layout([layout_lib.UNSHARDED, _MESH_DIM_X, _MESH_DIM_Y],

self.mesh)

b_layout = Layout([_MESH_DIM_X, layout_lib.UNSHARDED, layout_lib.UNSHARDED],

self.mesh)

# If any input is sharded on the concat dim, then the concat dim is

# replicated in the output. Dim 0 in the output is replicated because of

# broadcast compatibility, mesh dimension X is already used in dim 1 of

# input a.

output_layout = Layout(

[layout_lib.UNSHARDED, layout_lib.UNSHARDED, _MESH_DIM_Y], self.mesh)

a = api.relayout(a, a_layout)

b = api.relayout(b, b_layout)

@polymorphic_function.function

def concat_fn(a, b):

return op([a, b])