def _split_tensor(tensor, num_splits):
if tensor is None:
return [None] * num_splits
elif isinstance(tensor, sparse_tensor.SparseTensor):
return sparse_ops.sparse_split_v2(tensor, num_splits, axis=0)
else:
return array_ops.split(tensor, num_splits)
def _split_tensor(tensor, num_splits):
"""Splits tensor into num_splits pieces, returns a list of pieces."""
if tensor is None:
return [None] * num_splits
elif num_splits <= 0:
return ValueError(
'Tensors cannot be split into {} pieces.'.format(num_splits))
elif num_splits == 1:
return [tensor]
elif isinstance(tensor, sparse_tensor.SparseTensor):
return sparse_ops.sparse_split_v2(tensor, num_splits, axis=0)
else:
return array_ops.split(tensor, num_splits)
def _sparse_dense_matmul_deep_support(st,
dt,
st_shape=None,
dt_shape=None):
if st_shape is None:
st_shape = st.get_shape().as_list()
if len(st_shape) == 2:
if len(dt_shape) > 2:
return _sparse_dense_matmul_deep_support(st,
dt[0],
st_shape=st_shape,
dt_shape=dt_shape[1:])
return tf.sparse.sparse_dense_matmul(st, dt)
if dt_shape is None:
dt_shape = dt.get_shape().as_list()
st_split = sparse_split_v2(st, num_split=st_shape[0], axis=0)
res = []
n_shape = st_shape[1:]
n_dt_shape = dt_shape[1:]
if st_shape[0] != dt_shape[0]:
for i in range(st_shape[0]):
ss = st_split[i]
ss = tf.sparse.reshape(ss, shape=n_shape)
r = _sparse_dense_matmul_deep_support(ss,
dt,
st_shape=n_shape,
dt_shape=dt_shape)
res.append(tf.expand_dims(r, axis=0))
else:
dt_split = tf.split(dt, num_or_size_splits=dt_shape[0], axis=0)
for i in range(st_shape[0]):
ss = st_split[i]
ds = dt_split[i]
ss = tf.sparse.reshape(ss, shape=n_shape)
ds = tf.reshape(ds, shape=n_dt_shape)
r = _sparse_dense_matmul_deep_support(ss,
ds,
st_shape=n_shape,
dt_shape=n_dt_shape)
res.append(tf.expand_dims(r, axis=0))
return tf.concat(res, axis=0)
