tensorlayer3/tensorlayer/array_ops.py

#! /usr/bin/python
# -*- coding: utf-8 -*-
"""A file containing functions related to array manipulation."""

from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op, dtypes, ops, tensor_shape
from tensorflow.python.framework.constant_op import constant
from tensorflow.python.framework.ops import convert_to_tensor
from tensorflow.python.ops.array_ops import shape_internal
from tensorflow.python.ops.gen_array_ops import fill, reshape

__all__ = ['alphas', 'alphas_like']


def alphas(shape, alpha_value, name=None):
    """Creates a tensor with all elements set to `alpha_value`.
    This operation returns a tensor of type `dtype` with shape `shape` and all
    elements set to alpha.

    Parameters
    ----------
    shape: A list of integers, a tuple of integers, or a 1-D `Tensor` of type `int32`.
        The shape of the desired tensor
    alpha_value: `float32`, `float64`, `int8`, `uint8`, `int16`, `uint16`, int32`, `int64`
        The value used to fill the resulting `Tensor`.
    name: str
        A name for the operation (optional).

    Returns
    -------
    A `Tensor` with all elements set to alpha.

    Examples
    --------
    >>> tl.alphas([2, 3], tf.int32)  # [[alpha, alpha, alpha], [alpha, alpha, alpha]]
    """
    with ops.name_scope(name, "alphas", [shape]) as name:

        alpha_tensor = convert_to_tensor(alpha_value)
        alpha_dtype = dtypes.as_dtype(alpha_tensor.dtype).base_dtype

        if not isinstance(shape, ops.Tensor):
            try:
                shape = constant_op._tensor_shape_tensor_conversion_function(tensor_shape.TensorShape(shape))
            except (TypeError, ValueError):
                shape = ops.convert_to_tensor(shape, dtype=dtypes.int32)

        if not shape._shape_tuple():
            shape = reshape(shape, [-1])  # Ensure it's a vector

        try:
            output = constant(alpha_value, shape=shape, dtype=alpha_dtype, name=name)

        except (TypeError, ValueError):
            output = fill(shape, constant(alpha_value, dtype=alpha_dtype), name=name)

        if output.dtype.base_dtype != alpha_dtype:
            raise AssertionError("Dtypes do not corresponds: %s and %s" % (output.dtype.base_dtype, alpha_dtype))

        return output


def alphas_like(tensor, alpha_value, name=None, optimize=True):
    """Creates a tensor with all elements set to `alpha_value`.
    Given a single tensor (`tensor`), this operation returns a tensor of the same
    type and shape as `tensor` with all elements set to `alpha_value`.

    Parameters
    ----------
    tensor: tf.Tensor
        The Tensorflow Tensor that will be used as a template.
    alpha_value: `float32`, `float64`, `int8`, `uint8`, `int16`, `uint16`, int32`, `int64`
        The value used to fill the resulting `Tensor`.
    name: str
        A name for the operation (optional).
    optimize: bool
        if true, attempt to statically determine the shape of 'tensor' and encode it as a constant.

    Returns
    -------
    A `Tensor` with all elements set to `alpha_value`.

    Examples
    --------
    >>> tensor = tf.constant([[1, 2, 3], [4, 5, 6]])
    >>> tl.alphas_like(tensor, 0.5)  # [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]]
    """
    with ops.name_scope(name, "alphas_like", [tensor]) as name:
        tensor = ops.convert_to_tensor(tensor, name="tensor")

        if context.in_eager_mode():  # and dtype is not None and dtype != tensor.dtype:
            ret = alphas(shape_internal(tensor, optimize=optimize), alpha_value=alpha_value, name=name)

        else:  # if context.in_graph_mode():

            # For now, variant types must be created via zeros_like; as we need to
            # pass the input variant object to the proper zeros callback.

            if (optimize and tensor.shape.is_fully_defined()):
                # We can produce a zeros tensor independent of the value of 'tensor',
                # since the shape is known statically.
                ret = alphas(tensor.shape, alpha_value=alpha_value, name=name)

            # elif dtype is not None and dtype != tensor.dtype and dtype != dtypes.variant:
            else:
                ret = alphas(shape_internal(tensor, optimize=optimize), alpha_value=alpha_value, name=name)

            ret.set_shape(tensor.get_shape())

        return ret