wehub-resource-sync
332 lines
9.5 KiB
Python
332 lines
9.5 KiB
Python
# Copyright 2023 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.
|
|
# ==============================================================================
|
|
|
|
"""DTensor helpers for random generators."""
|
|
|
|
from tensorflow.dtensor.python import api
|
|
from tensorflow.dtensor.python import layout as layout_lib
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import tensor_shape
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import gen_math_ops
|
|
from tensorflow.python.ops import gen_stateless_random_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import shape_util
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# stateless rngs
|
|
# ------------------------------------------------------------------------------
|
|
|
|
|
|
# TODO(b/171746536): switch all rng ops to official versions once supported.
|
|
def _old_tf_random_stateless_normal(
|
|
shape,
|
|
seed,
|
|
mean=0.0,
|
|
stddev=1.0,
|
|
dtype=dtypes.float32,
|
|
name=None,
|
|
layout=None,
|
|
):
|
|
"""DTensor stateless normal implementation that takes an layout."""
|
|
with ops.name_scope(
|
|
name, "stateless_random_normal", [shape, seed, mean, stddev]
|
|
) as name:
|
|
seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed")
|
|
shape = shape_util.shape_tensor(shape)
|
|
mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
|
|
stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
|
|
rnd = api.call_with_layout(
|
|
gen_stateless_random_ops.stateless_random_normal,
|
|
layout,
|
|
shape,
|
|
seed,
|
|
dtype,
|
|
)
|
|
result = math_ops.add(rnd * stddev, mean, name=name)
|
|
shape_util.maybe_set_static_shape(result, shape)
|
|
return result
|
|
|
|
|
|
def _old_tf_random_stateless_uniform(
|
|
shape,
|
|
seed,
|
|
minval=0,
|
|
maxval=None,
|
|
dtype=dtypes.float32,
|
|
name=None,
|
|
layout=None,
|
|
):
|
|
"""DTensor stateless uniform implementation that takes an layout."""
|
|
dtype = dtypes.as_dtype(dtype)
|
|
accepted_dtypes = (
|
|
dtypes.float16,
|
|
dtypes.bfloat16,
|
|
dtypes.float32,
|
|
dtypes.float64,
|
|
dtypes.int32,
|
|
dtypes.int64,
|
|
dtypes.uint32,
|
|
dtypes.uint64,
|
|
)
|
|
if dtype not in accepted_dtypes:
|
|
raise ValueError(
|
|
f"Argument `dtype` got invalid value {dtype}. Accepted dtypes are "
|
|
f"{accepted_dtypes}."
|
|
)
|
|
if dtype.is_integer:
|
|
if (minval is None) != (maxval is None):
|
|
raise ValueError(
|
|
f"For integer `dtype` argument {dtype}, argument `minval` and "
|
|
f"`maxval` must be both None or not None. Got `minval`={minval} and "
|
|
f"`maxval`={maxval}."
|
|
)
|
|
if minval is not None and dtype in (dtypes.uint32, dtypes.uint64):
|
|
raise ValueError(
|
|
f"Argument `dtype` got invalid value {dtype} when argument `minval` "
|
|
"is not None. Please don't use unsigned integers in this case."
|
|
)
|
|
|
|
shape = shape_util.shape_tensor(shape)
|
|
with ops.name_scope(
|
|
name, "stateless_random_uniform", [shape, seed, minval, maxval]
|
|
) as name:
|
|
seed = ops.convert_to_tensor(seed, dtype_hint=dtypes.int32, name="seed")
|
|
|
|
if dtype.is_integer and minval is None and maxval is None:
|
|
result = api.call_with_layout(
|
|
gen_stateless_random_ops.stateless_random_uniform_full_int,
|
|
layout,
|
|
shape,
|
|
seed=seed,
|
|
dtype=dtype,
|
|
name=name,
|
|
)
|
|
else:
|
|
if not dtype.is_integer and maxval is None:
|
|
maxval = 1
|
|
val_range = ops.convert_to_tensor(
|
|
maxval - minval, dtype=dtype, name="range"
|
|
)
|
|
minval = ops.convert_to_tensor(minval, dtype=dtype, name="min")
|
|
if dtype.is_integer:
|
|
result = api.call_with_layout(
|
|
gen_stateless_random_ops.stateless_random_uniform_int,
|
|
layout,
|
|
shape,
|
|
seed=seed,
|
|
minval=minval,
|
|
maxval=maxval,
|
|
)
|
|
else:
|
|
rnd = api.call_with_layout(
|
|
gen_stateless_random_ops.stateless_random_uniform,
|
|
layout,
|
|
shape,
|
|
seed=seed,
|
|
dtype=dtype,
|
|
)
|
|
result = math_ops.add(rnd * val_range, minval, name=name)
|
|
shape_util.maybe_set_static_shape(result, shape)
|
|
return result
|
|
|
|
|
|
def _old_tf_stateless_truncated_normal(
|
|
shape,
|
|
seed,
|
|
mean=0.0,
|
|
stddev=1.0,
|
|
dtype=dtypes.float32,
|
|
name=None,
|
|
layout=None,
|
|
):
|
|
"""DTensor stateless truncated normal implementation that takes an layout."""
|
|
with ops.name_scope(
|
|
name, "stateless_truncated_normal", [shape, seed, mean, stddev]
|
|
) as name:
|
|
seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed")
|
|
shape = shape_util.shape_tensor(shape)
|
|
mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean")
|
|
stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev")
|
|
rnd = api.call_with_layout(
|
|
gen_stateless_random_ops.stateless_truncated_normal,
|
|
layout,
|
|
shape,
|
|
seed,
|
|
dtype,
|
|
)
|
|
result = math_ops.add(rnd * stddev, mean, name=name)
|
|
shape_util.maybe_set_static_shape(result, shape)
|
|
return result
|
|
|
|
|
|
def stateless_random_normal(
|
|
shape,
|
|
seed,
|
|
mean=0.0,
|
|
stddev=1.0,
|
|
dtype=dtypes.float32,
|
|
name=None,
|
|
layout=None,
|
|
):
|
|
"""DTensor stateless RNG."""
|
|
if not context.executing_eagerly():
|
|
layout = None
|
|
|
|
return _old_tf_random_stateless_normal(
|
|
shape,
|
|
seed=seed,
|
|
mean=mean,
|
|
stddev=stddev,
|
|
dtype=dtype,
|
|
name=name,
|
|
layout=layout,
|
|
)
|
|
|
|
|
|
def stateless_random_uniform(
|
|
shape,
|
|
seed,
|
|
minval=0,
|
|
maxval=None,
|
|
dtype=dtypes.float32,
|
|
name=None,
|
|
layout=None,
|
|
):
|
|
"""DTensor stateless random uniform."""
|
|
if not context.executing_eagerly():
|
|
layout = None
|
|
|
|
return _old_tf_random_stateless_uniform(
|
|
shape,
|
|
seed=seed,
|
|
minval=minval,
|
|
maxval=maxval,
|
|
dtype=dtype,
|
|
name=name,
|
|
layout=layout,
|
|
)
|
|
|
|
|
|
def stateless_truncated_normal(
|
|
shape,
|
|
seed,
|
|
mean=0.0,
|
|
stddev=1.0,
|
|
dtype=dtypes.float32,
|
|
name=None,
|
|
layout=None,
|
|
):
|
|
"""DTensor stateless RNG."""
|
|
if not context.executing_eagerly():
|
|
layout = None
|
|
|
|
return _old_tf_stateless_truncated_normal(
|
|
shape,
|
|
seed=seed,
|
|
mean=mean,
|
|
stddev=stddev,
|
|
dtype=dtype,
|
|
name=name,
|
|
layout=layout,
|
|
)
|
|
|
|
|
|
def stateless_split(seed, num=2, mesh=None):
|
|
seed = ops.convert_to_tensor(seed)
|
|
layout = None
|
|
if mesh:
|
|
layout = layout_lib.Layout.replicated(mesh, rank=2)
|
|
return stateless_random_uniform(
|
|
shape=[num, 2],
|
|
seed=seed,
|
|
dtype=seed.dtype,
|
|
minval=None,
|
|
maxval=None,
|
|
layout=layout,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# stateless dropout.
|
|
# ------------------------------------------------------------------------------
|
|
|
|
|
|
def _get_noise_shape(x, noise_shape):
|
|
"""Noisve shape util copied from tf nn_ops."""
|
|
# If noise_shape is none return immediately.
|
|
if noise_shape is None:
|
|
return array_ops.shape(x)
|
|
|
|
try:
|
|
# Best effort to figure out the intended shape.
|
|
# If not possible, let the op to handle it.
|
|
# In eager mode exception will show up.
|
|
noise_shape_ = tensor_shape.as_shape(noise_shape)
|
|
except (TypeError, ValueError):
|
|
return noise_shape
|
|
|
|
if x.shape.dims is not None and len(x.shape.dims) == len(noise_shape_.dims):
|
|
new_dims = []
|
|
for i, dim in enumerate(x.shape.dims):
|
|
if noise_shape_.dims[i].value is None and dim.value is not None:
|
|
new_dims.append(dim.value)
|
|
else:
|
|
new_dims.append(noise_shape_.dims[i].value)
|
|
return tensor_shape.TensorShape(new_dims)
|
|
|
|
return noise_shape
|
|
|
|
|
|
# TODO(b/171213877, b/169909066): Fix layout prop in function case for the rng
|
|
# Op used. The layout prop should be able to propagate the layout from input
|
|
# tensor `x` to the tf.mul and then back propagate the layout to the
|
|
# `random_tensor`.
|
|
def dropout(x, rate, noise_shape=None, seed=None, name=None):
|
|
"""DTensor replacement for dropout."""
|
|
if not isinstance(rate, float):
|
|
raise ValueError("rate should be float for dropout.")
|
|
if seed is None:
|
|
raise ValueError("seed must be specified for DTensor dropout. Got: None")
|
|
|
|
with ops.name_scope(name, "dropout", [x]):
|
|
x_dtype = x.dtype
|
|
keep_prob = 1 - rate
|
|
scale = 1 / keep_prob
|
|
scale = ops.convert_to_tensor(scale, dtype=x_dtype)
|
|
ret = gen_math_ops.mul(x, scale)
|
|
|
|
noise_shape = _get_noise_shape(x, noise_shape)
|
|
# stateless_random_uniform requires a shape [2] seed.
|
|
seed = [seed, 0]
|
|
|
|
if context.executing_eagerly():
|
|
layout = api.fetch_layout(x)
|
|
else:
|
|
layout = None
|
|
random_tensor = _old_tf_random_stateless_uniform(
|
|
noise_shape, seed=seed, minval=0, maxval=1, dtype=x_dtype, layout=layout
|
|
)
|
|
keep_mask = random_tensor >= rate
|
|
ret = gen_math_ops.mul(ret, gen_math_ops.cast(keep_mask, x_dtype))
|
|
if not context.executing_eagerly():
|
|
ret.set_shape(x.get_shape())
|
|
return ret
|
|
|
|
|
|
# TODO(b/195413777): error out for stateful dropout.
|