Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
chunks = [chunkss[i] for i in out_ind]
if adjust_chunks:
for i, ind in enumerate(out_ind):
if ind in adjust_chunks:
if callable(adjust_chunks[ind]):
chunks[i] = tuple(map(adjust_chunks[ind], chunks[i]))
elif isinstance(adjust_chunks[ind], int):
chunks[i] = tuple(adjust_chunks[ind] for _ in chunks[i])
elif isinstance(adjust_chunks[ind], (tuple, list)):
chunks[i] = tuple(adjust_chunks[ind])
else:
raise NotImplementedError(
"adjust_chunks values must be callable, int, or tuple")
chunks = tuple(chunks)
return Array(sharedict.merge((out, dsk), *dsks), out, chunks, dtype=dtype)
}
dtype = np.bincount([]).dtype
dsk[(final_name, 0)] = (_bincount_sum, list(dsk), dtype)
graph = HighLevelGraph.from_collections(
final_name, dsk, dependencies=[x] if weights is None else [x, weights]
)
if minlength == 0:
chunks = ((np.nan,),)
else:
chunks = ((minlength,),)
meta = meta_from_array(x, 1, dtype=dtype)
return Array(graph, final_name, chunks, meta=meta)
old = x.name
keys = list(product(*map(range, x.numblocks)))
offsets = list(product(*(accumulate(operator.add, bd[:-1], 0) for bd in x.chunks)))
if ravel:
offset_info = zip(offsets, repeat(x.shape))
else:
offset_info = pluck(axis[0], offsets)
chunks = tuple((1,) * len(c) if i in axis else c for (i, c) in enumerate(x.chunks))
dsk = dict(
((name,) + k, (chunk, (old,) + k, axis, off))
for (k, off) in zip(keys, offset_info)
)
# The dtype of `tmp` doesn't actually matter, just need to provide something
graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x])
tmp = Array(graph, name, chunks, dtype=x.dtype)
dtype = np.argmin([1]).dtype
result = _tree_reduce(tmp, agg, axis, False, dtype, split_every, combine)
return handle_out(out, result)
def _take_dask_array_from_numpy(a, indices, axis):
assert isinstance(a, np.ndarray)
assert isinstance(indices, Array)
return indices.map_blocks(
lambda block: np.take(a, block, axis), chunks=indices.chunks, dtype=a.dtype
)
uu, ss, vvh = np.linalg.svd(np.ones(shape=(1, 1), dtype=data.dtype))
k = _nanmin(m, n) # avoid RuntimeWarning with np.nanmin([m, n])
m_u = m
n_u = int(k) if not np.isnan(k) else k
n_s = n_u
m_vh = n_u
n_vh = n
d_vh = max(m_vh, n_vh) # full matrix returned: but basically n
graph = HighLevelGraph(layers, dependencies)
u_meta = meta_from_array(data, len((m_u, n_u)), uu.dtype)
s_meta = meta_from_array(data, len((n_s,)), ss.dtype)
vh_meta = meta_from_array(data, len((d_vh, d_vh)), vvh.dtype)
u = Array(
graph,
name_u_st4,
shape=(m_u, n_u),
chunks=(data.chunks[0], (n_u,)),
meta=u_meta,
)
s = Array(graph, name_s_st2, shape=(n_s,), chunks=((n_s,),), meta=s_meta)
vh = Array(
graph, name_v_st2, shape=(d_vh, d_vh), chunks=((n,), (n,)), meta=vh_meta
)
return u, s, vh
leaf_arrs = []
for i, (ocd, oax, meta) in enumerate(zip(output_coredimss, output_axes, metas)):
core_output_shape = tuple(core_shapes[d] for d in ocd)
core_chunkinds = len(ocd) * (0,)
output_shape = loop_output_shape + core_output_shape
output_chunks = loop_output_chunks + core_output_shape
leaf_name = "%s_%d-%s" % (name, i, token)
leaf_dsk = {
(leaf_name,)
+ key[1:]
+ core_chunkinds: ((getitem, key, i) if nout else key)
for key in keys
}
graph = HighLevelGraph.from_collections(leaf_name, leaf_dsk, dependencies=[tmp])
meta = meta_from_array(meta, len(output_shape))
leaf_arr = Array(
graph, leaf_name, chunks=output_chunks, shape=output_shape, meta=meta
)
### Axes:
if keepdims:
slices = len(leaf_arr.shape) * (slice(None),) + len(oax) * (np.newaxis,)
leaf_arr = leaf_arr[slices]
tidcs = [None] * len(leaf_arr.shape)
for i, oa in zip(range(-len(oax), 0), oax):
tidcs[oa] = i
j = 0
for i in range(len(tidcs)):
if tidcs[i] is None:
tidcs[i] = j
j += 1
@dispatch(TensorDot, Array, Array)
def compute_up(expr, lhs, rhs, **kwargs):
left_index = list(alphabet[:ndim(lhs)])
right_index = list(ALPHABET[:ndim(rhs)])
out_index = left_index + right_index
for l, r in zip(expr._left_axes, expr._right_axes):
out_index.remove(right_index[r])
out_index.remove(left_index[l])
right_index[r] = left_index[l]
func = many(binop=np.tensordot, reduction=sum,
axes=(expr._left_axes, expr._right_axes))
return atop(func,
next(names), out_index,
lhs, tuple(left_index),
rhs, tuple(right_index))
# some functions can't compute empty arrays (those for which reduced_meta
# fall into the ValueError exception) and we have to rely on reshaping
# the array according to len(out_chunks)
if is_arraylike(meta) and meta.ndim != len(out_chunks):
if len(out_chunks) == 0:
meta = meta.sum()
else:
meta = meta.reshape((0,) * len(out_chunks))
if np.isscalar(meta):
return Array(graph, name, out_chunks, dtype=dtype)
else:
with ignoring(AttributeError):
meta = meta.astype(dtype)
return Array(graph, name, out_chunks, meta=meta)
)
for i, sub_block_info in enumerate(all_blocks)
)
layers[name_r_stacked] = dsk_r_stacked
dependencies[name_r_stacked] = {name_r_st1}
# retrieve R_stacked for recursion with tsqr
vchunks_rstacked = tuple(
[sum(map(lambda x: x[1], sub_block_info)) for sub_block_info in all_blocks]
)
graph = HighLevelGraph(layers, dependencies)
# dsk.dependencies[name_r_stacked] = {data.name}
r_stacked_meta = meta_from_array(
data, len((sum(vchunks_rstacked), n)), dtype=rr.dtype
)
r_stacked = Array(
graph,
name_r_stacked,
shape=(sum(vchunks_rstacked), n),
chunks=(vchunks_rstacked, n),
meta=r_stacked_meta,
)
# recurse
q_inner, r_inner = tsqr(r_stacked, _max_vchunk_size=cr_max)
layers = toolz.merge(q_inner.dask.layers, r_inner.dask.layers)
dependencies = toolz.merge(q_inner.dask.dependencies, r_inner.dask.dependencies)
# Q_inner: "unstack"
name_q_st2 = "getitem" + token + "-q2"
dsk_q_st2 = dict(
(
if j > 0:
prevs = []
for p in range(j):
prev = name_lt_dot, j, p, i, p
dsk[prev] = (np.dot, (name, j, p), (name_upper, p, i))
prevs.append(prev)
target = (operator.sub, target, (sum, prevs))
dsk[name_upper, j, i] = (_solve_triangular_lower, (name, j, j), target)
dsk[name, i, j] = (np.transpose, (name_upper, j, i))
graph_upper = HighLevelGraph.from_collections(name_upper, dsk, dependencies=[a])
graph_lower = HighLevelGraph.from_collections(name, dsk, dependencies=[a])
cho = scipy.linalg.cholesky(np.array([[1, 2], [2, 5]], dtype=a.dtype))
meta = meta_from_array(a, dtype=cho.dtype)
lower = Array(graph_lower, name, shape=a.shape, chunks=a.chunks, meta=meta)
# do not use .T, because part of transposed blocks are already calculated
upper = Array(graph_upper, name_upper, shape=a.shape, chunks=a.chunks, meta=meta)
return lower, upper