8aa9b39c709e1eb42b6ac81aac39c0b2b451c670,mesh_tensorflow/layers.py,,dense,#Any#Any#Any#Any#Any#Any#Any#Any#Any#Any#,27

Before Change

  
  if variable_dtype is None:
    variable_dtype = mtf.VariableDType(master_dtype, slice_dtype, x.dtype)
  if variable_dtype.activation_dtype != x.dtype:
    raise ValueError("variable_dtype.activation_dtype must match x.dtype "
                     "variable_dtype=%s x=%s" % (variable_dtype, x))
  if expert_dims is None:
    expert_dims = []
  if reduced_dims is None:
    reduced_dims = x.shape.dims[-1:]

After Change

        w_shape,
        initializer=tf.random_normal_initializer(stddev=stddev),
        dtype=variable_dtype)
    w = mtf.cast(w, x.dtype)
    y = mtf.einsum([x, w], output_shape)
    if use_bias:
      b = mtf.get_variable(
          x.mesh,

Instances