jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        fun: Callable[[TensorType], TensorType],
        lb: Sequence[float] | TensorType,
        ub: Sequence[float] | TensorType,
        keep_feasible: bool = False,
    ):
        # Implement caching to avoid calling the constraint function multiple times to get value
        # and gradient.
        def _constraint_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType]:
            val, grad = tfp.math.value_and_gradient(fun, x)

            tf.debugging.assert_shapes(
                [(val, [..., 1])],
                message="Nonlinear constraint only supports single output function.",
            )

            return t.cast(val, dtype=x.dtype), t.cast(grad, dtype=x.dtype)

        cache_x: TensorType = tf.constant([])
        cache_f: TensorType = tf.constant([])
        cache_df_dx: TensorType = tf.constant([])

        def val_fun(x: TensorType) -> TensorType:
            nonlocal cache_x, cache_f, cache_df_dx
            if not np.array_equal(x, cache_x):
                cache_f, cache_df_dx = _constraint_value_and_gradient(x)
                cache_x = x
            return cache_f

        def jac_fun(x: TensorType) -> TensorType:
            nonlocal cache_x, cache_f, cache_df_dx
            if not np.array_equal(x, cache_x):
                cache_f, cache_df_dx = _constraint_value_and_gradient(x)
                cache_x = x
            return cache_df_dx

        self._orig_fun = fun  # Used for constraints equality check.
        super().__init__(val_fun, lb, ub, jac=jac_fun, keep_feasible=keep_feasible)
//  [python]
def __init__(
        self,
        fun: Callable[[TensorType], TensorType],
        lb: Sequence[float] | TensorType,
        ub: Sequence[float] | TensorType,
        keep_feasible: bool = False,
    ):
        # Implement caching to avoid calling the constraint function multiple times to get value
        # and gradient.
        def _constraint_value_and_gradient(x: TensorType) -> Tuple[TensorType, TensorType]:
            val, grad = tfp.math.value_and_gradient(fun, x)

            tf.debugging.assert_shapes(
                [(val, [..., 1])],
                message="Nonlinear constraint only supports single output function.",
            )

            return tf.cast(val, dtype=x.dtype), tf.cast(grad, dtype=x.dtype)

        cache_x: TensorType = tf.constant([])
        cache_f: TensorType = tf.constant([])
        cache_df_dx: TensorType = tf.constant([])

        def val_fun(x: TensorType) -> TensorType:
            nonlocal cache_x, cache_f, cache_df_dx
            if not np.array_equal(x, cache_x):
                cache_f, cache_df_dx = _constraint_value_and_gradient(x)
                cache_x = x
            return cache_f

        def jac_fun(x: TensorType) -> TensorType:
            nonlocal cache_x, cache_f, cache_df_dx
            if not np.array_equal(x, cache_x):
                cache_f, cache_df_dx = _constraint_value_and_gradient(x)
                cache_x = x
            return cache_df_dx

        self._orig_fun = fun  # Used for constraints equality check.
        super().__init__(val_fun, lb, ub, jac=jac_fun, keep_feasible=keep_feasible)
//  [python]
def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] = None):
        r"""
        Build a :class:`TaggedProductSearchSpace` from a list ``spaces`` of other spaces. If
        ``tags`` are provided then they form the identifiers of the subspaces, otherwise the
        subspaces are labelled numerically.

        :param spaces: A sequence of :class:`SearchSpace` objects representing the space's subspaces
        :param tags: An optional list of tags giving the unique identifiers of
            the space's subspaces.
        :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different
            length to ``tags`` when ``tags`` is provided or if ``tags`` contains duplicates.
        """

        super().__init__(spaces, tags)
        subspace_sizes = self.subspace_dimension

        self._subspace_sizes_by_tag = dict(zip(self._tags, subspace_sizes))

        self._subspace_starting_indices = dict(
            zip(self._tags, t.cumsum(subspace_sizes, exclusive=True))
        )

        self._dimension = t.reduce_sum(subspace_sizes)
//  [python]
def __init__(self, spaces: Sequence[SearchSpace], tags: Optional[Sequence[str]] = None):
        r"""
        Build a :class:`TaggedProductSearchSpace` from a list ``spaces`` of other spaces. If
        ``tags`` are provided then they form the identifiers of the subspaces, otherwise the
        subspaces are labelled numerically.

        :param spaces: A sequence of :class:`SearchSpace` objects representing the space's subspaces
        :param tags: An optional list of tags giving the unique identifiers of
            the space's subspaces.
        :raise ValueError (or tf.errors.InvalidArgumentError): If ``spaces`` has a different
            length to ``tags`` when ``tags`` is provided or if ``tags`` contains duplicates.
        """

        super().__init__(spaces, tags)
        subspace_sizes = self.subspace_dimension

        self._subspace_sizes_by_tag = dict(zip(self._tags, subspace_sizes))

        self._subspace_starting_indices = dict(
            zip(self._tags, tf.cumsum(subspace_sizes, exclusive=True))
        )

        self._dimension = tf.cast(tf.reduce_sum(subspace_sizes), dtype=tf.int32)
//  [python]
def call(self, inputs, training=True):
        if len(inputs) != 2:
            raise ValueError("Must have two in_layers")
        mean_parent, std_parent = inputs[0], inputs[1]
        noise_scale = tf.cast(training or not self.training_only, tf.float32)
        from tensorflow.python.ops import array_ops
        sample_noise = tf.random.normal(array_ops.shape(mean_parent),
                                        0,
                                        self.noise_epsilon,
                                        dtype=tf.float32)
        return mean_parent + noise_scale * std_parent * sample_noise
//  [python]
def get_cells(self):
        """Returns the locations of all grid points in box.

        Suppose start is -10 Angstrom, stop is 10 Angstrom, nbr_cutoff is 1.
        Then would return a list of length 20^3 whose entries would be
        [(-10, -10, -10), (-10, -10, -9), ..., (9, 9, 9)]

        Returns
        -------
        cells: tf.Tensor
            (n_cells, ndim) shape.
        """
        start, stop, nbr_cutoff = self.start, self.stop, self.nbr_cutoff
        mesh_args = [
            tf.range(start, stop, nbr_cutoff) for _ in range(self.ndim)
        ]
        return t.cast(
            tf.reshape(tf.transpose(tf.stack(tf.meshgrid(*mesh_args))),
                       (self.n_cells, self.ndim)), tf.float32)
//  [python]
def get_cells(self):
        """Returns the locations of all grid points in box.

        Suppose start is -10 Angstrom, stop is 10 Angstrom, nbr_cutoff is 1.
        Then would return a list of length 20^3 whose entries would be
        [(-10, -10, -10), (-10, -10, -9), ..., (9, 9, 9)]

        Returns
        -------
        cells: tf.Tensor
            (n_cells, ndim) shape.
        """
        start, stop, nbr_cutoff = self.start, self.stop, self.nbr_cutoff
        mesh_args = [
            tf.range(start, stop, nbr_cutoff) for _ in range(self.ndim)
        ]
        return tf.cast(
            tf.reshape(tf.transpose(tf.stack(tf.meshgrid(*mesh_args))),
                       (self.n_cells, self.ndim)), tf.float32)
//  [python]
def call(self, inputs):
        """In layers should be of shape dtype tf.float32, (None, self.max_atoms, 4)"""
        atom_numbers = t.cast(inputs[:, :, 0], t.int32)
        flags = t.sign(atom_numbers)
        flags = t.cast(
            t.expand_dims(flags, 1) * t.expand_dims(flags, 2), t.float32)
        coordinates = inputs[:, :, 1:]
        if self.coordinates_in_bohr:
            coordinates = coordinates * 0.52917721092

        d = self.distance_matrix(coordinates, flags)

        d_radial_cutoff = self.distance_cutoff(d, self.radial_cutoff, flags)
        d_angular_cutoff = self.distance_cutoff(d, self.angular_cutoff, flags)

        radial_sym = self.radial_symmetry(d_radial_cutoff, d, atom_numbers)
        angular_sym = self.angular_symmetry(d_angular_cutoff, d, atom_numbers,
                                            coordinates)
        return t.concat([
            t.cast(t.expand_dims(atom_numbers, 2), t.float32), radial_sym,
            angular_sym
        ],
                         axis=2)
//  [python]
def call(self, inputs):
        """In layers should be of shape dtype tf.float32, (None, self.max_atoms, 4)"""
        atom_numbers = tf.cast(inputs[:, :, 0], tf.int32)
        flags = tf.sign(atom_numbers)
        flags = tf.cast(
            tf.expand_dims(flags, 1) * tf.expand_dims(flags, 2), tf.float32)
        coordinates = inputs[:, :, 1:]
        if self.coordinates_in_bohr:
            coordinates = coordinates * 0.52917721092

        d = self.distance_matrix(coordinates, flags)

        d_radial_cutoff = self.distance_cutoff(d, self.radial_cutoff, flags)
        d_angular_cutoff = self.distance_cutoff(d, self.angular_cutoff, flags)

        radial_sym = self.radial_symmetry(d_radial_cutoff, d, atom_numbers)
        angular_sym = self.angular_symmetry(d_angular_cutoff, d, atom_numbers,
                                            coordinates)
        return tf.concat([
            tf.cast(tf.expand_dims(atom_numbers, 2), tf.float32), radial_sym,
            angular_sym
        ],
                         axis=2)
//  [python]
def radial_symmetry(self, d_cutoff, d, atom_numbers):
        """ Radial Symmetry Function """
        embedding = tf.eye(np.max(self.atom_cases) + 1)
        atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers)

        Rs = np.linspace(0., self.radial_cutoff, self.radial_length)
        ita = np.ones_like(Rs) * 3 / (Rs[1] - Rs[0])**2
        Rs = t.cast(np.reshape(Rs, (1, 1, 1, -1)), tf.float32)
        ita = t.cast(np.reshape(ita, (1, 1, 1, -1)), tf.float32)
        length = ita.get_shape().as_list()[-1]

        d_cutoff = tf.stack([d_cutoff] * length, axis=3)
        d = tf.stack([d] * length, axis=3)

        out = tf.exp(-ita * tf.square(d - Rs)) * d_cutoff
        if self.atomic_number_differentiated:
            out_tensors = []
            for atom_type in self.atom_cases:
                selected_atoms = tf.expand_dims(tf.expand_dims(
                    atom_numbers_embedded[:, :, atom_type], axis=1),
                                                axis=3)
                out_tensors.append(tf.reduce_sum(out * selected_atoms, axis=2))
            return tf.concat(out_tensors, axis=2)
        else:
            return tf.reduce_sum(out, axis=2)
//  [python]
def radial_symmetry(self, d_cutoff, d, atom_numbers):
        """ Radial Symmetry Function """
        embedding = tf.eye(np.max(self.atom_cases) + 1)
        atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers)

        Rs = np.linspace(0., self.radial_cutoff, self.radial_length)
        ita = np.ones_like(Rs) * 3 / (Rs[1] - Rs[0])**2
        Rs = tf.cast(np.reshape(Rs, (1, 1, 1, -1)), tf.float32)
        ita = tf.cast(np.reshape(ita, (1, 1, 1, -1)), tf.float32)
        length = ita.get_shape().as_list()[-1]

        d_cutoff = tf.stack([d_cutoff] * length, axis=3)
        d = tf.stack([d] * length, axis=3)

        out = tf.exp(-ita * tf.square(d - Rs)) * d_cutoff
        if self.atomic_number_differentiated:
            out_tensors = []
            for atom_type in self.atom_cases:
                selected_atoms = tf.expand_dims(tf.expand_dims(
                    atom_numbers_embedded[:, :, atom_type], axis=1),
                                                axis=3)
                out_tensors.append(tf.reduce_sum(out * selected_atoms, axis=2))
            return tf.concat(out_tensors, axis=2)
        else:
            return tf.reduce_sum(out, axis=2)
//  [python]
def position_embedding_sine(attention_mask,
                            num_pos_features=256,
                            temperature=10000.,
                            normalize=True,
                            scale=2 * math.pi):
  """Sine-based positional embeddings for 2D images.

  Args:
    attention_mask: a `bool` Tensor specifying the size of the input image to
      the Transformer and which elements are padded, of size [batch_size,
      height, width]
    num_pos_features: a `int` specifying the number of positional features,
      should be equal to the hidden size of the Transformer network
    temperature: a `float` specifying the temperature of the positional
      embedding. Any type that is converted to a `float` can also be accepted.
    normalize: a `bool` determining whether the positional embeddings should be
      normalized between [0, scale] before application of the sine and cos
      functions.
    scale: a `float` if normalize is True specifying the scale embeddings before
      application of the embedding function.

  Returns:
    embeddings: a `float` tensor of the same shape as input_tensor specifying
      the positional embeddings based on sine features.
  """
  if num_pos_features % 2 != 0:
    raise ValueError(
        "Number of embedding features (num_pos_features) must be even when "
        "column and row embeddings are concatenated.")
  num_pos_features = num_pos_features // 2

  # Produce row and column embeddings based on total size of the image
  # <tf.float>[batch_size, height, width]
  attention_mask = tf.cast(attention_mask, tf.float32)
  row_embedding = tf.cumsum(attention_mask, 1)
  col_embedding = tf.cumsum(attention_mask, 2)

  if normalize:
    eps = 1e-6
    row_embedding = row_embedding / (row_embedding[:, -1:, :] + eps) * scale
    col_embedding = col_embedding / (col_embedding[:, :, -1:] + eps) * scale

  dim_t = tf.range(num_pos_features, dtype=row_embedding.dtype)
  dim_t = tf.pow(temperature, 2 * (dim_t // 2) / num_pos_features)

  # Creates positional embeddings for each row and column position
  # <tf.float>[batch_size, height, width, num_pos_features]
  pos_row = tf.expand_dims(row_embedding, -1) / dim_t
  pos_col = tf.expand_dims(col_embedding, -1) / dim_t
  pos_row = tf.stack(
      [tf.sin(pos_row[:, :, :, 0::2]),
       tf.cos(pos_row[:, :, :, 1::2])], axis=4)
  pos_col = tf.stack(
      [tf.sin(pos_col[:, :, :, 0::2]),
       tf.cos(pos_col[:, :, :, 1::2])], axis=4)

  # final_shape = pos_row.shape.as_list()[:3] + [-1]
  final_shape = tf_utils.get_shape_list(pos_row)[:3] + [-1]
  pos_row = tf.reshape(pos_row, final_shape)
  pos_col = tf.reshape(pos_col, final_shape)
  output = tf.concat([pos_row, pos_col], -1)

  embeddings = tf.cast(output, tf.float32)
  return embeddings
//  [python]
def compute_nbr_list(self, coords):
        """Get closest neighbors for atoms.

        Needs to handle padding for atoms with no neighbors.

        Parameters
        ----------
        coords: tf.Tensor
            Shape (N_atoms, ndim)

        Returns
        -------
        nbr_list: tf.Tensor
            Shape (N_atoms, M_nbrs) of atom indices
        """
        # Shape (n_cells, ndim)
        cells = self.get_cells()

        # List of length N_atoms, each element of different length uniques_i
        nbrs = self.get_atoms_in_nbrs(coords, cells)
        padding = tf.fill((self.M_nbrs,), -1)
        padded_nbrs = [
            tf.concat([unique_nbrs, padding], 0) for unique_nbrs in nbrs
        ]

        # List of length N_atoms, each element of different length uniques_i
        # List of length N_atoms, each a tensor of shape
        # (uniques_i, ndim)
        nbr_coords = [tf.gather(coords, atom_nbrs) for atom_nbrs in nbrs]

        # Add phantom atoms that exist far outside the box
        coord_padding = tf.fill((self.M_nbrs, self.ndim), 2 * self.stop)
        padded_nbr_coords = [
            tf.concat([nbr_coord, coord_padding], 0) for nbr_coord in nbr_coords
        ]

        # List of length N_atoms, each of shape (1, ndim)
        atom_coords = tf.split(coords, self.N_atoms)
        # TODO(rbharath): How does distance need to be modified here to
        # account for periodic boundary conditions?
        # List of length N_atoms each of shape (M_nbrs)
        padded_dists = [
            tf.reduce_sum((atom_coord - padded_nbr_coord)**2, axis=1)
            for (atom_coord,
                 padded_nbr_coord) in zip(atom_coords, padded_nbr_coords)
        ]

        padded_closest_nbrs = [
            tf.nn.top_k(-padded_dist, k=self.M_nbrs)[1]
            for padded_dist in padded_dists
        ]

        # N_atoms elts of size (M_nbrs,) each
        padded_neighbor_list = [
            tf.gather(padded_atom_nbrs, padded_closest_nbr)
            for (padded_atom_nbrs,
                 padded_closest_nbr) in zip(padded_nbrs, padded_closest_nbrs)
        ]

        neighbor_list = tf.stack(padded_neighbor_list)

        return neighbor_list
//  [python]
def compute_nbr_list(self, coords):
        """Get closest neighbors for atoms.

        Needs to handle padding for atoms with no neighbors.

        Parameters
        ----------
        coords: tf.Tensor
            Shape (N_atoms, ndim)

        Returns
        -------
        nbr_list: tf.Tensor
            Shape (N_atoms, M_nbrs) of atom indices
        """
        # Shape (n_cells, ndim)
        cells = self.get_cells()

        # List of length N_atoms, each element of different length uniques_i
        nbrs = self.get_atoms_in_nbrs(coords, cells)
        padding = tf.fill((self.M_nbrs,), -1)
        padded_nbrs = [
            tf.concat([unique_nbrs, padding], 0) for unique_nbrs in nbrs
        ]

        # List of length N_atoms, each element of different length uniques_i
        # List of length N_atoms, each a tensor of shape
        # (uniques_i, ndim)
        nbr_coords = [tf.gather(coords, atom_nbrs) for atom_nbrs in nbrs]

        # Add phantom atoms that exist far outside the box
        coord_padding = tf.cast(
            tf.fill((self.M_nbrs, self.ndim), 2 * self.stop), tf.float32)
        padded_nbr_coords = [
            tf.concat([nbr_coord, coord_padding], 0) for nbr_coord in nbr_coords
        ]

        # List of length N_atoms, each of shape (1, ndim)
        atom_coords = tf.split(coords, self.N_atoms)
        # TODO(rbharath): How does distance need to be modified here to
        # account for periodic boundary conditions?
        # List of length N_atoms each of shape (M_nbrs)
        padded_dists = [
            tf.reduce_sum((atom_coord - padded_nbr_coord)**2, axis=1)
            for (atom_coord,
                 padded_nbr_coord) in zip(atom_coords, padded_nbr_coords)
        ]

        padded_closest_nbrs = [
            tf.nn.top_k(-padded_dist, k=self.M_nbrs)[1]
            for padded_dist in padded_dists
        ]

        # N_atoms elts of size (M_nbrs,) each
        padded_neighbor_list = [
            tf.gather(padded_atom_nbrs, padded_closest_nbr)
            for (padded_atom_nbrs,
                 padded_closest_nbr) in zip(padded_nbrs, padded_closest_nbrs)
        ]

        neighbor_list = tf.stack(padded_neighbor_list)

        return neighbor_list
//  [python]
def call(self, inputs):
        """
        Parameters
        ----------
        X: tf.Tensor of shape (B, N, d)
            Coordinates/features.
        Nbrs: tf.Tensor of shape (B, N, M)
            Neighbor list.
        Nbrs_Z: tf.Tensor of shape (B, N, M)
            Atomic numbers of neighbor atoms.

        Returns
        -------
        layer: tf.Tensor of shape (B, N, l)
            A new tensor representing the output of the atomic conv layer
        """
        X = inputs[0]
        Nbrs = t.cast(inputs[1], t.int32)
        Nbrs_Z = inputs[2]

        # N: Maximum number of atoms
        # M: Maximum number of neighbors
        # d: Number of coordinates/features/filters
        # B: Batch Size
        N = X.get_shape()[-2]
        d = X.get_shape()[-1]
        M = Nbrs.get_shape()[-1]
        B = X.get_shape()[0]

        # Compute the distances and radial symmetry functions.
        D = self.distance_tensor(X, Nbrs, self.boxsize, B, N, M, d)
        R = self.distance_matrix(D)
        R = t.expand_dims(R, 0)
        rsf = self.radial_symmetry_function(R, self.rc, self.rs, self.re)

        if not self.atom_types:
            cond = t.cast(t.not_equal(Nbrs_Z, 0), t.float32)
            cond = t.reshape(cond, (1, -1, N, M))
            layer = t.reduce_sum(cond * rsf, 3)
        else:
            sym = []
            for j in range(len(self.atom_types)):
                cond = t.cast(t.equal(Nbrs_Z, self.atom_types[j]), t.float32)
                cond = t.reshape(cond, (1, -1, N, M))
                sym.append(t.reduce_sum(cond * rsf, 3))
            layer = t.concat(sym, 0)

        layer = t.transpose(layer, [1, 2, 0])  # (l, B, N) -> (B, N, l)
        m, v = t.nn.moments(layer, axes=[0])

        return t.nn.batch_normalization(layer, m, v, None, None, 1e-3)
//  [python]
def call(self, inputs):
        """
        Parameters
        ----------
        X: tf.Tensor of shape (B, N, d)
            Coordinates/features.
        Nbrs: tf.Tensor of shape (B, N, M)
            Neighbor list.
        Nbrs_Z: tf.Tensor of shape (B, N, M)
            Atomic numbers of neighbor atoms.

        Returns
        -------
        layer: tf.Tensor of shape (B, N, l)
            A new tensor representing the output of the atomic conv layer
        """
        X = inputs[0]
        Nbrs = tf.cast(inputs[1], tf.int32)
        Nbrs_Z = inputs[2]

        # N: Maximum number of atoms
        # M: Maximum number of neighbors
        # d: Number of coordinates/features/filters
        # B: Batch Size
        N = X.get_shape()[-2]
        d = X.get_shape()[-1]
        M = Nbrs.get_shape()[-1]
        B = X.get_shape()[0]

        # Compute the distances and radial symmetry functions.
        D = self.distance_tensor(X, Nbrs, self.boxsize, B, N, M, d)
        R = self.distance_matrix(D)
        R = tf.expand_dims(R, 0)
        rsf = self.radial_symmetry_function(R, self.rc, self.rs, self.re)

        if not self.atom_types:
            cond = tf.cast(tf.not_equal(Nbrs_Z, 0), tf.float32)
            cond = tf.reshape(cond, (1, -1, N, M))
            layer = tf.reduce_sum(cond * rsf, 3)
        else:
            sym = []
            for j in range(len(self.atom_types)):
                cond = tf.cast(tf.equal(Nbrs_Z, self.atom_types[j]), tf.float32)
                cond = tf.reshape(cond, (1, -1, N, M))
                sym.append(tf.reduce_sum(cond * rsf, 3))
            layer = tf.concat(sym, 0)

        layer = tf.transpose(layer, [1, 2, 0])  # (l, B, N) -> (B, N, l)
        m, v = tf.nn.moments(layer, axes=[0])

        return tf.nn.batch_normalization(layer, m, v, None, None, 1e-3)
//  [python]
def result(self):
        p, q, _, _ = self._compute_variables()
        n = t.cast(self.n, t.float32)
        non_zeros_col = t.math.count_nonzero(self.ncol)
        non_zeros_row = t.math.count_nonzero(self.nrow)
        m = t.cast(t.minimum(non_zeros_col, non_zeros_row), t.float32)
        return 2 * (p - q) / (t.square(n) * (m - 1) / m)
//  [python]
def result(self):
        p, q, _, _ = self._compute_variables()
        n = tf.cast(self.n, tf.float32)
        non_zeros_col = tf.math.count_nonzero(self.ncol)
        non_zeros_row = tf.math.count_nonzero(self.nrow)
        m = tf.cast(tf.minimum(non_zeros_col, non_zeros_row), tf.float32)
        return 2 * (p - q) / (tf.square(n) * (m - 1) / m)
//  [python]
def result(self):
        nrow = t.cast(self.nrow, t.float32)
        ncol = t.cast(self.ncol, t.float32)
        n = t.cast(self.n, t.float32)

        nrow_ = tf.where(nrow > 0, nrow, -1.0)
        rrow = tf.cumsum(nrow, exclusive=True) + (nrow_ - n) / 2
        ncol_ = tf.where(ncol > 0, ncol, -1.0)
        rcol = tf.cumsum(ncol, exclusive=True) + (ncol_ - n) / 2

        rrow = rrow / tf.math.sqrt(tf.reduce_sum(nrow * tf.square(rrow)))
        rcol = rcol / tf.math.sqrt(tf.reduce_sum(ncol * tf.square(rcol)))

        m = t.cast(self.m, t.float32)
        corr = tf.matmul(tf.expand_dims(rrow, axis=0), m)
        corr = tf.matmul(corr, tf.expand_dims(rcol, axis=1))
        return tf.squeeze(corr)
//  [python]
def result(self):
        nrow = tf.cast(self.nrow, tf.float32)
        ncol = tf.cast(self.ncol, tf.float32)
        n = tf.cast(self.n, tf.float32)

        nrow_ = tf.where(nrow > 0, nrow, -1.0)
        rrow = tf.cumsum(nrow, exclusive=True) + (nrow_ - n) / 2
        ncol_ = tf.where(ncol > 0, ncol, -1.0)
        rcol = tf.cumsum(ncol, exclusive=True) + (ncol_ - n) / 2

        rrow = rrow / tf.math.sqrt(tf.reduce_sum(nrow * tf.square(rrow)))
        rcol = rcol / tf.math.sqrt(tf.reduce_sum(ncol * tf.square(rcol)))

        m = tf.cast(self.m, tf.float32)
        corr = tf.matmul(tf.expand_dims(rrow, axis=0), m)
        corr = tf.matmul(corr, tf.expand_dims(rcol, axis=1))
        return tf.squeeze(corr)
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):
        """Updates `m`, `nrow`, `ncol` respectively the joint and
        marginal histograms of (`y_true`, `y_pred`)
        """

        y_true = tf.clip_by_value(y_true, self.actual_min, self.actual_max)
        y_pred = tf.clip_by_value(y_pred, self.preds_min, self.preds_max)

        i = (
            tf.searchsorted(
                self.actual_cuts,
                tf.reshape(y_true, [-1]),
                side="right",
                out_type=tf.int64,
            )
            - 1
        )
        j = (
            tf.searchsorted(
                self.preds_cuts,
                tf.reshape(y_pred, [-1]),
                side="right",
                out_type=tf.int64,
            )
            - 1
        )

        nrow = tf.tensor_scatter_nd_add(
            self.nrow, tf.expand_dims(i, axis=-1), tf.ones_like(i)
        )
        ncol = tf.tensor_scatter_nd_add(
            self.ncol, tf.expand_dims(j, axis=-1), tf.ones_like(j)
        )
        ij = tf.stack([i, j], axis=1)
        m = tf.tensor_scatter_nd_add(self.m, ij, tf.ones_like(i))

        self.n.assign_add(tf.shape(i, out_type=tf.int64)[0])
        self.m.assign(m)
        self.nrow.assign(nrow)
        self.ncol.assign(ncol)
//  [python]
def update_state(self, y_true, y_pred, sample_weight=None):
        """Updates `m`, `nrow`, `ncol` respectively the joint and
        marginal histograms of (`y_true`, `y_pred`)
        """

        y_true = tf.clip_by_value(y_true, self.actual_min, self.actual_max)
        y_pred = tf.clip_by_value(y_pred, self.preds_min, self.preds_max)

        i = (
            tf.searchsorted(
                self.actual_cuts,
                tf.cast(tf.reshape(y_true, [-1]), self.actual_cuts.dtype),
                side="right",
                out_type=tf.int64,
            )
            - 1
        )
        j = (
            tf.searchsorted(
                self.preds_cuts,
                tf.cast(tf.reshape(y_pred, [-1]), self.preds_cuts.dtype),
                side="right",
                out_type=tf.int64,
            )
            - 1
        )

        nrow = tf.tensor_scatter_nd_add(
            self.nrow, tf.expand_dims(i, axis=-1), tf.ones_like(i)
        )
        ncol = tf.tensor_scatter_nd_add(
            self.ncol, tf.expand_dims(j, axis=-1), tf.ones_like(j)
        )
        ij = tf.stack([i, j], axis=1)
        m = tf.tensor_scatter_nd_add(self.m, ij, tf.ones_like(i))

        self.n.assign_add(tf.shape(i, out_type=tf.int64)[0])
        self.m.assign(m)
        self.nrow.assign(nrow)
        self.ncol.assign(ncol)
//  [python]
def _compute_variables(self):
        """Compute a tuple containing the concordant pairs, discordant pairs,
        ties in `y_true` and `y_pred`.

        Returns:
          A tuple
        """
        m = t.cast(self.m, t.float32)
        n_cap = t.cumsum(t.cumsum(m, axis=0), axis=1)
        # Number of concordant pairs.
        p = t.math.reduce_sum(t.multiply(n_cap[:-1, :-1], m[1:, 1:]))
        sum_m_squard = t.math.reduce_sum(t.math.square(m))
        # Ties in x.
        t = (
            t.cast(t.math.reduce_sum(t.math.square(self.nrow)), t.float32)
            - sum_m_squard
        ) / 2.0
        # Ties in y.
        u = (
            t.cast(t.math.reduce_sum(t.math.square(self.ncol)), t.float32)
            - sum_m_squard
        ) / 2.0
        # Ties in both.
        b = t.math.reduce_sum(t.multiply(m, (m - 1.0))) / 2.0
        # Number of discordant pairs.
        n = t.cast(self.n, t.float32)
        q = (n - 1.0) * n / 2.0 - p - t - u - b
        return p, q, t, u
//  [python]
def _compute_variables(self):
        """Compute a tuple containing the concordant pairs, discordant pairs,
        ties in `y_true` and `y_pred`.

        Returns:
          A tuple
        """
        m = tf.cast(self.m, tf.float32)
        n_cap = tf.cumsum(tf.cumsum(m, axis=0), axis=1)
        # Number of concordant pairs.
        p = tf.math.reduce_sum(tf.multiply(n_cap[:-1, :-1], m[1:, 1:]))
        sum_m_squard = tf.math.reduce_sum(tf.math.square(m))
        # Ties in x.
        t = (
            tf.cast(tf.math.reduce_sum(tf.math.square(self.nrow)), tf.float32)
            - sum_m_squard
        ) / 2.0
        # Ties in y.
        u = (
            tf.cast(tf.math.reduce_sum(tf.math.square(self.ncol)), tf.float32)
            - sum_m_squard
        ) / 2.0
        # Ties in both.
        b = tf.math.reduce_sum(tf.multiply(m, (m - 1.0))) / 2.0
        # Number of discordant pairs.
        n = tf.cast(self.n, tf.float32)
        q = (n - 1.0) * n / 2.0 - p - t - u - b
        return p, q, t, u
//  [python]
def call(self, inputs, training=True):
        """
        parent layers: atom_features, parents, calculation_orders, calculation_masks, n_atoms
        """
        atom_features = inputs[0]
        # each atom corresponds to a graph, which is represented by the `max_atoms*max_atoms` int32 matrix of index
        # each gragh include `max_atoms` of steps(corresponding to rows) of calculating graph features
        parents = t.cast(inputs[1], dtype=t.int32)
        # target atoms for each step: (batch_size*max_atoms) * max_atoms
        calculation_orders = inputs[2]
        calculation_masks = inputs[3]

        n_atoms = tf.squeeze(inputs[4])
        graph_features = tf.zeros((self.max_atoms * self.batch_size,
                                   self.max_atoms + 1, self.n_graph_feat))

        for count in range(self.max_atoms):
            # `count`-th step
            # extracting atom features of target atoms: (batch_size*max_atoms) * n_atom_features
            mask = calculation_masks[:, count]
            current_round = tf.boolean_mask(calculation_orders[:, count], mask)
            batch_atom_features = tf.gather(atom_features, current_round)

            # generating index for graph features used in the inputs
            stack1 = tf.reshape(
                tf.stack([tf.boolean_mask(tf.range(n_atoms), mask)] *
                         (self.max_atoms - 1),
                         axis=1), [-1])
            stack2 = tf.reshape(tf.boolean_mask(parents[:, count, 1:], mask),
                                [-1])
            index = tf.stack([stack1, stack2], axis=1)
            # extracting graph features for parents of the target atoms, then flatten
            # shape: (batch_size*max_atoms) * [(max_atoms-1)*n_graph_features]
            batch_graph_features = tf.reshape(
                tf.gather_nd(graph_features, index),
                [-1, (self.max_atoms - 1) * self.n_graph_feat])

            # concat into the input tensor: (batch_size*max_atoms) * n_inputs
            batch_inputs = tf.concat(
                axis=1, values=[batch_atom_features, batch_graph_features])
            # DAGgraph_step maps from batch_inputs to a batch of graph_features
            # of shape: (batch_size*max_atoms) * n_graph_features
            # representing the graph features of target atoms in each graph
            batch_outputs = _DAGgraph_step(batch_inputs, self.W_list,
                                           self.b_list, self.activation_fn,
                                           self.dropouts, training)

            # index for targe atoms
            target_index = tf.stack([tf.range(n_atoms), parents[:, count, 0]],
                                    axis=1)
            target_index = tf.boolean_mask(target_index, mask)
            graph_features = tf.tensor_scatter_nd_update(
                graph_features, target_index, batch_outputs)
        return batch_outputs
//  [python]
def call(self, inputs, training=True):
        """
        parent layers: atom_features, parents, calculation_orders, calculation_masks, n_atoms
        """
        atom_features = inputs[0]
        # each atom corresponds to a graph, which is represented by the `max_atoms*max_atoms` int32 matrix of index
        # each gragh include `max_atoms` of steps(corresponding to rows) of calculating graph features
        parents = tf.cast(inputs[1], dtype=tf.int32)
        # target atoms for each step: (batch_size*max_atoms) * max_atoms
        calculation_orders = inputs[2]
        calculation_masks = inputs[3]

        n_atoms = tf.squeeze(inputs[4])
        graph_features = tf.zeros((self.max_atoms * self.batch_size,
                                   self.max_atoms + 1, self.n_graph_feat))

        for count in range(self.max_atoms):
            # `count`-th step
            # extracting atom features of target atoms: (batch_size*max_atoms) * n_atom_features
            mask = calculation_masks[:, count]
            current_round = tf.boolean_mask(calculation_orders[:, count], mask)
            batch_atom_features = tf.gather(atom_features, current_round)

            # generating index for graph features used in the inputs
            stack1 = tf.reshape(
                tf.stack([tf.boolean_mask(tf.range(n_atoms), mask)] *
                         (self.max_atoms - 1),
                         axis=1), [-1])
            stack2 = tf.reshape(tf.boolean_mask(parents[:, count, 1:], mask),
                                [-1])
            index = tf.stack([stack1, stack2], axis=1)
            # extracting graph features for parents of the target atoms, then flatten
            # shape: (batch_size*max_atoms) * [(max_atoms-1)*n_graph_features]
            batch_graph_features = tf.reshape(
                tf.gather_nd(graph_features, index),
                [-1, (self.max_atoms - 1) * self.n_graph_feat])

            # concat into the input tensor: (batch_size*max_atoms) * n_inputs
            batch_inputs = tf.concat(
                axis=1, values=[batch_atom_features, batch_graph_features])
            # DAGgraph_step maps from batch_inputs to a batch of graph_features
            # of shape: (batch_size*max_atoms) * n_graph_features
            # representing the graph features of target atoms in each graph
            batch_outputs = _DAGgraph_step(batch_inputs, self.W_list,
                                           self.b_list, self.activation_fn,
                                           self.dropouts, training)

            # index for targe atoms
            target_index = tf.stack([tf.range(n_atoms), parents[:, count, 0]],
                                    axis=1)
            target_index = tf.boolean_mask(target_index, mask)
            graph_features = tf.tensor_scatter_nd_update(
                graph_features, target_index, batch_outputs)
        return batch_outputs
//  [python]
def __init__(
        self,
        actual_min: float = 0.0,
        actual_max: float = 1.0,
        preds_min: float = 0.0,
        preds_max: float = 1.0,
        actual_cutpoints: int = 100,
        preds_cutpoints: int = 100,
        name: str = None,
        dtype: AcceptableDTypes = None,
    ):
        """Creates a `CorrelationBase` instance.

        Args:
          actual_min: the inclusive lower bound on values from actual.
          actual_max: the exclusive upper bound on values from actual.
          preds_min: the inclusive lower bound on values from preds.
          preds_max: the exclusive upper bound on values from preds.
          actual_cutpoints: the number of divisions to create in actual range,
            defaults to 100.
          preds_cutpoints: the number of divisions to create in preds range,
            defaults to 100.
          name: (optional) String name of the metric instance
          dtype: (optional) Data type of the metric result. Defaults to `None`
        """
        super().__init__(name=name, dtype=dtype)
        self.actual_min = actual_min
        self.actual_max = actual_max
        self.preds_min = preds_min
        self.preds_max = preds_max
        self.actual_cutpoints = actual_cutpoints
        self.preds_cutpoints = preds_cutpoints
        actual_cuts = np.linspace(
            t.cast(self.actual_min, t.float32),
            t.cast(self.actual_max, t.float32),
            self.actual_cutpoints,
        )
        actual_cuts[-1] += backend.epsilon()
        preds_cuts = np.linspace(
            t.cast(self.preds_min, t.float32),
            t.cast(self.preds_max, t.float32),
            self.preds_cutpoints,
        )
        preds_cuts[-1] += backend.epsilon()
        self.actual_cuts = tf.convert_to_tensor(actual_cuts, tf.float32)
        self.preds_cuts = tf.convert_to_tensor(preds_cuts, tf.float32)
        self.m = self.add_weight(
            "m", (self.actual_cutpoints - 1, self.preds_cutpoints - 1), dtype=tf.int64
        )
        self.nrow = self.add_weight("nrow", (self.actual_cutpoints - 1), dtype=tf.int64)
        self.ncol = self.add_weight("ncol", (self.preds_cutpoints - 1), dtype=tf.int64)
        self.n = self.add_weight("n", (), dtype=tf.int64)
//  [python]
def __init__(
        self,
        actual_min: float = 0.0,
        actual_max: float = 1.0,
        preds_min: float = 0.0,
        preds_max: float = 1.0,
        actual_cutpoints: int = 100,
        preds_cutpoints: int = 100,
        name: str = None,
        dtype: AcceptableDTypes = None,
    ):
        """Creates a `CorrelationBase` instance.

        Args:
          actual_min: the inclusive lower bound on values from actual.
          actual_max: the exclusive upper bound on values from actual.
          preds_min: the inclusive lower bound on values from preds.
          preds_max: the exclusive upper bound on values from preds.
          actual_cutpoints: the number of divisions to create in actual range,
            defaults to 100.
          preds_cutpoints: the number of divisions to create in preds range,
            defaults to 100.
          name: (optional) String name of the metric instance
          dtype: (optional) Data type of the metric result. Defaults to `None`
        """
        super().__init__(name=name, dtype=dtype)
        self.actual_min = actual_min
        self.actual_max = actual_max
        self.preds_min = preds_min
        self.preds_max = preds_max
        self.actual_cutpoints = actual_cutpoints
        self.preds_cutpoints = preds_cutpoints
        actual_cuts = np.linspace(
            tf.cast(self.actual_min, tf.float32),
            tf.cast(self.actual_max, tf.float32),
            self.actual_cutpoints,
        )
        actual_cuts[-1] += backend.epsilon()
        preds_cuts = np.linspace(
            tf.cast(self.preds_min, tf.float32),
            tf.cast(self.preds_max, tf.float32),
            self.preds_cutpoints,
        )
        preds_cuts[-1] += backend.epsilon()
        self.actual_cuts = tf.convert_to_tensor(actual_cuts, tf.float32)
        self.preds_cuts = tf.convert_to_tensor(preds_cuts, tf.float32)
        self.m = self.add_weight(
            "m", (self.actual_cutpoints - 1, self.preds_cutpoints - 1), dtype=tf.int64
        )
        self.nrow = self.add_weight("nrow", (self.actual_cutpoints - 1), dtype=tf.int64)
        self.ncol = self.add_weight("ncol", (self.preds_cutpoints - 1), dtype=tf.int64)
        self.n = self.add_weight("n", (), dtype=tf.int64)
//  [python]
def angular_symmetry(self, d_cutoff, d, atom_numbers, coordinates):
        """ Angular Symmetry Function """

        max_atoms = self.max_atoms
        embedding = tf.eye(np.max(self.atom_cases) + 1)
        atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers)

        Rs = np.linspace(0., self.angular_cutoff, self.angular_length)
        ita = 3 / (Rs[1] - Rs[0])**2
        thetas = np.linspace(0., np.pi, self.angular_length)
        zeta = float(self.angular_length**2)

        ita, zeta, Rs, thetas = np.meshgrid(ita, zeta, Rs, thetas)
        zeta = t.cast(np.reshape(zeta, (1, 1, 1, 1, -1)), tf.float32)
        ita = t.cast(np.reshape(ita, (1, 1, 1, 1, -1)), tf.float32)
        Rs = t.cast(np.reshape(Rs, (1, 1, 1, 1, -1)), tf.float32)
        thetas = t.cast(np.reshape(thetas, (1, 1, 1, 1, -1)), tf.float32)
        length = zeta.get_shape().as_list()[-1]

        # tf.stack issues again...
        vector_distances = tf.stack([coordinates] * max_atoms, 1) - tf.stack(
            [coordinates] * max_atoms, 2)
        R_ij = tf.stack([d] * max_atoms, axis=3)
        R_ik = tf.stack([d] * max_atoms, axis=2)
        f_R_ij = tf.stack([d_cutoff] * max_atoms, axis=3)
        f_R_ik = tf.stack([d_cutoff] * max_atoms, axis=2)

        # Define angle theta = arccos(R_ij(Vector) dot R_ik(Vector)/R_ij(distance)/R_ik(distance))
        vector_mul = tf.reduce_sum(tf.stack([vector_distances] * max_atoms, axis=3) * \
                                   tf.stack([vector_distances] * max_atoms, axis=2), axis=4)
        vector_mul = vector_mul * tf.sign(f_R_ij) * tf.sign(f_R_ik)
        theta = tf.acos(tf.math.divide(vector_mul, R_ij * R_ik + 1e-5))

        R_ij = tf.stack([R_ij] * length, axis=4)
        R_ik = tf.stack([R_ik] * length, axis=4)
        f_R_ij = tf.stack([f_R_ij] * length, axis=4)
        f_R_ik = tf.stack([f_R_ik] * length, axis=4)
        theta = tf.stack([theta] * length, axis=4)

        out_tensor = tf.pow((1. + tf.cos(theta - thetas)) / 2., zeta) * \
                     tf.exp(-ita * tf.square((R_ij + R_ik) / 2. - Rs)) * f_R_ij * f_R_ik * 2

        if self.atomic_number_differentiated:
            out_tensors = []
            for id_j, atom_type_j in enumerate(self.atom_cases):
                for atom_type_k in self.atom_cases[id_j:]:
                    selected_atoms = tf.stack([atom_numbers_embedded[:, :, atom_type_j]] * max_atoms, axis=2) * \
                                     tf.stack([atom_numbers_embedded[:, :, atom_type_k]] * max_atoms, axis=1)
                    selected_atoms = tf.expand_dims(tf.expand_dims(
                        selected_atoms, axis=1),
                                                    axis=4)
                    out_tensors.append(
                        tf.reduce_sum(out_tensor * selected_atoms, axis=(2, 3)))
            return tf.concat(out_tensors, axis=2)
        else:
            return tf.reduce_sum(out_tensor, axis=(2, 3))
//  [python]
def angular_symmetry(self, d_cutoff, d, atom_numbers, coordinates):
        """ Angular Symmetry Function """

        max_atoms = self.max_atoms
        embedding = tf.eye(np.max(self.atom_cases) + 1)
        atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers)

        Rs = np.linspace(0., self.angular_cutoff, self.angular_length)
        ita = 3 / (Rs[1] - Rs[0])**2
        thetas = np.linspace(0., np.pi, self.angular_length)
        zeta = float(self.angular_length**2)

        ita, zeta, Rs, thetas = np.meshgrid(ita, zeta, Rs, thetas)
        zeta = tf.cast(np.reshape(zeta, (1, 1, 1, 1, -1)), tf.float32)
        ita = tf.cast(np.reshape(ita, (1, 1, 1, 1, -1)), tf.float32)
        Rs = tf.cast(np.reshape(Rs, (1, 1, 1, 1, -1)), tf.float32)
        thetas = tf.cast(np.reshape(thetas, (1, 1, 1, 1, -1)), tf.float32)
        length = zeta.get_shape().as_list()[-1]

        # tf.stack issues again...
        vector_distances = tf.stack([coordinates] * max_atoms, 1) - tf.stack(
            [coordinates] * max_atoms, 2)
        R_ij = tf.stack([d] * max_atoms, axis=3)
        R_ik = tf.stack([d] * max_atoms, axis=2)
        f_R_ij = tf.stack([d_cutoff] * max_atoms, axis=3)
        f_R_ik = tf.stack([d_cutoff] * max_atoms, axis=2)

        # Define angle theta = arccos(R_ij(Vector) dot R_ik(Vector)/R_ij(distance)/R_ik(distance))
        vector_mul = tf.reduce_sum(tf.stack([vector_distances] * max_atoms, axis=3) * \
                                   tf.stack([vector_distances] * max_atoms, axis=2), axis=4)
        vector_mul = vector_mul * tf.sign(f_R_ij) * tf.sign(f_R_ik)
        theta = tf.acos(tf.math.divide(vector_mul, R_ij * R_ik + 1e-5))

        R_ij = tf.stack([R_ij] * length, axis=4)
        R_ik = tf.stack([R_ik] * length, axis=4)
        f_R_ij = tf.stack([f_R_ij] * length, axis=4)
        f_R_ik = tf.stack([f_R_ik] * length, axis=4)
        theta = tf.stack([theta] * length, axis=4)

        out_tensor = tf.pow((1. + tf.cos(theta - thetas)) / 2., zeta) * \
                     tf.exp(-ita * tf.square((R_ij + R_ik) / 2. - Rs)) * f_R_ij * f_R_ik * 2

        if self.atomic_number_differentiated:
            out_tensors = []
            for id_j, atom_type_j in enumerate(self.atom_cases):
                for atom_type_k in self.atom_cases[id_j:]:
                    selected_atoms = tf.stack([atom_numbers_embedded[:, :, atom_type_j]] * max_atoms, axis=2) * \
                                     tf.stack([atom_numbers_embedded[:, :, atom_type_k]] * max_atoms, axis=1)
                    selected_atoms = tf.expand_dims(tf.expand_dims(
                        selected_atoms, axis=1),
                                                    axis=4)
                    out_tensors.append(
                        tf.reduce_sum(out_tensor * selected_atoms, axis=(2, 3)))
            return tf.concat(out_tensors, axis=2)
        else:
            return tf.reduce_sum(out_tensor, axis=(2, 3))
//  [python]
def result(self):
        ncol = t.cast(self.ncol, t.float32)
        nrow = t.cast(self.nrow, t.float32)
        m = t.cast(self.m, t.float32)
        n = t.cast(self.n, t.float32)

        col_bins = (self.preds_cuts[1:] - self.preds_cuts[:-1]) / 2.0 + self.preds_cuts[
            :-1
        ]
        row_bins = (
            self.actual_cuts[1:] - self.actual_cuts[:-1]
        ) / 2.0 + self.actual_cuts[:-1]

        n_col = t.reduce_sum(ncol)
        n_row = t.reduce_sum(nrow)
        col_mean = t.reduce_sum(ncol * col_bins) / n_col
        row_mean = t.reduce_sum(nrow * row_bins) / n_row

        col_var = t.reduce_sum(ncol * t.square(col_bins)) - n_col * t.square(
            col_mean
        )
        row_var = t.reduce_sum(nrow * t.square(row_bins)) - n_row * t.square(
            row_mean
        )

        joint_product = m * t.expand_dims(row_bins, axis=1) * col_bins

        corr = t.reduce_sum(joint_product) - n * col_mean * row_mean

        return corr / t.sqrt(col_var * row_var)
//  [python]
def result(self):
        ncol = tf.cast(self.ncol, tf.float32)
        nrow = tf.cast(self.nrow, tf.float32)
        m = tf.cast(self.m, tf.float32)
        n = tf.cast(self.n, tf.float32)

        col_bins = (self.preds_cuts[1:] - self.preds_cuts[:-1]) / 2.0 + self.preds_cuts[
            :-1
        ]
        row_bins = (
            self.actual_cuts[1:] - self.actual_cuts[:-1]
        ) / 2.0 + self.actual_cuts[:-1]

        n_col = tf.reduce_sum(ncol)
        n_row = tf.reduce_sum(nrow)
        col_mean = tf.reduce_sum(ncol * col_bins) / n_col
        row_mean = tf.reduce_sum(nrow * row_bins) / n_row

        col_var = tf.reduce_sum(ncol * tf.square(col_bins)) - n_col * tf.square(
            col_mean
        )
        row_var = tf.reduce_sum(nrow * tf.square(row_bins)) - n_row * tf.square(
            row_mean
        )

        joint_product = m * tf.expand_dims(row_bins, axis=1) * col_bins

        corr = tf.reduce_sum(joint_product) - n * col_mean * row_mean

        return corr / tf.sqrt(col_var * row_var)
//  [python]
def switch(condition, then_expression, else_expression):
    '''Switches between two operations depending on a scalar value (int or bool).
    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    # Arguments
        condition: scalar tensor.
        then_expression: TensorFlow operation.
        else_expression: TensorFlow operation.
    '''
    x_shape = copy.copy(then_expression.get_shape())
    x = tf.cond(tf.cast(condition, 'bool'),
                lambda: then_expression,
                lambda: else_expression)
    x.set_shape(x_shape)
    return x