jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def lsuv_init(model, batch, verbose=True, margin=0.1, max_iter=100):
    """
    Initialize neural network weights using layer-sequential unit-variance
    initialization.

    See:
        Mishkin and Matas, "All you need is a good init". 2016.
        https://arxiv.org/abs/1511.06422

    Parameters
    ----------
    model : keras.Model
    batch : dict
        Training data, as would be passed keras.Model.fit()
    verbose : boolean
        Whether to print progress to stdout
    margin : float
    max_iter : int

    Returns
    -------
    keras.Model
        Same as what was passed in.
    """
    configure_tensorflow()
    from tensorflow.keras.layers import Dense, Convolution2D
    needed_variance = 1.0
    layers_inintialized = 0
    for layer in model.layers:
        if not isinstance(layer, (Dense, Convolution2D)):
            continue
        # avoid small layers where activation variance close to zero, esp.
        # for small batches_generator
        if n.prod(layer.get_output_shape_at(0)[1:]) < 32:
            if verbose:
                print('LSUV initialization skipping', layer.name)
            continue
        layers_inintialized += 1
        weights_and_biases = layer.get_weights()
        weights_and_biases[0] = svd_orthonormal(weights_and_biases[0].shape)
        layer.set_weights(weights_and_biases)
        activations = get_activations(model, layer, batch)
        variance = numpy.var(activations)
        iteration = 0
        if verbose:
            print(layer.name, variance)
        while abs(needed_variance - variance) > margin:
            if verbose:
                print(
                    'LSUV initialization',
                    layer.name,
                    iteration,
                    needed_variance,
                    margin,
                    variance)

            if numpy.abs(numpy.sqrt(variance)) < 1e-7:
                break  # avoid zero division

            weights_and_biases = layer.get_weights()
            weights_and_biases[0] /= numpy.sqrt(variance) / numpy.sqrt(
                needed_variance)
            layer.set_weights(weights_and_biases)
            activations = get_activations(model, layer, batch)
            variance = numpy.var(activations)

            iteration += 1
            if iteration >= max_iter:
                break
    if verbose:
        print('Done with LSUV: total layers initialized', layers_inintialized)
    return model
//  [python]
def lsuv_init(model, batch, verbose=True, margin=0.1, max_iter=100):
    """
    Initialize neural network weights using layer-sequential unit-variance
    initialization.

    See:
        Mishkin and Matas, "All you need is a good init". 2016.
        https://arxiv.org/abs/1511.06422

    Parameters
    ----------
    model : keras.Model
    batch : dict
        Training data, as would be passed keras.Model.fit()
    verbose : boolean
        Whether to print progress to stdout
    margin : float
    max_iter : int

    Returns
    -------
    keras.Model
        Same as what was passed in.
    """
    configure_tensorflow()
    from tensorflow.keras.layers import Dense, Convolution2D
    needed_variance = 1.0
    layers_inintialized = 0
    for layer in model.layers:
        if not isinstance(layer, (Dense, Convolution2D)):
            continue
        # avoid small layers where activation variance close to zero, esp.
        # for small batches_generator
        if numpy.prod(layer.get_output_shape_at(0)[1:]) < 32:
            if verbose:
                print('LSUV initialization skipping', layer.name)
            continue
        layers_inintialized += 1
        weights_and_biases = layer.get_weights()
        weights_and_biases[0] = svd_orthonormal(weights_and_biases[0].shape)
        layer.set_weights(weights_and_biases)
        activations = get_activations(model, layer, batch)
        variance = numpy.var(activations)
        iteration = 0
        if verbose:
            print(layer.name, variance)
        while abs(needed_variance - variance) > margin:
            if verbose:
                print(
                    'LSUV initialization',
                    layer.name,
                    iteration,
                    needed_variance,
                    margin,
                    variance)

            if numpy.abs(numpy.sqrt(variance)) < 1e-7:
                break  # avoid zero division

            weights_and_biases = layer.get_weights()
            weights_and_biases[0] /= numpy.sqrt(variance) / numpy.sqrt(
                needed_variance)
            layer.set_weights(weights_and_biases)
            activations = get_activations(model, layer, batch)
            variance = numpy.var(activations)

            iteration += 1
            if iteration >= max_iter:
                break
    if verbose:
        print('Done with LSUV: total layers initialized', layers_inintialized)
    return model
//  [python]
def test_advinc_subtensor1():
    """ Test the second case in the opt local_gpu_advanced_incsubtensor1 """
    for shp in [(3, 3), (3, 3, 3)]:
        shared = cuda.shared_constructor
        xval = numpy.arange(n.prod(shp), dtype='float32').reshape(shp) + 1
        yval = numpy.empty((2,) + shp[1:], dtype='float32')
        yval[:] = 10
        x = shared(xval, name='x')
        y = T.tensor(dtype='float32',
                     broadcastable=(False,) * len(shp),
                     name='y')
        expr = T.advanced_inc_subtensor1(x, y, [0, 2])
        f = theano.function([y], expr, mode=mode_with_gpu)
        assert sum([isinstance(node.op, cuda.GpuAdvancedIncSubtensor1)
                    for node in f.maker.fgraph.toposort()]) == 1
        rval = f(yval)
        rep = xval.copy()
        rep[[0, 2]] += yval
        utt.assert_allclose(rval, rep)
//  [python]
def test_advinc_subtensor1():
    """ Test the second case in the opt local_gpu_advanced_incsubtensor1 """
    for shp in [(3, 3), (3, 3, 3)]:
        shared = cuda.shared_constructor
        xval = numpy.arange(numpy.prod(shp), dtype='float32').reshape(shp) + 1
        yval = numpy.empty((2,) + shp[1:], dtype='float32')
        yval[:] = 10
        x = shared(xval, name='x')
        y = T.tensor(dtype='float32',
                     broadcastable=(False,) * len(shp),
                     name='y')
        expr = T.advanced_inc_subtensor1(x, y, [0, 2])
        f = theano.function([y], expr, mode=mode_with_gpu)
        assert sum([isinstance(node.op, cuda.GpuAdvancedIncSubtensor1)
                    for node in f.maker.fgraph.toposort()]) == 1
        rval = f(yval)
        rep = xval.copy()
        rep[[0, 2]] += yval
        utt.assert_allclose(rval, rep)
//  [python]
def _isNonAxi(Pot):
    """
    Determine whether this potential is non-axisymmetric

    Parameters
    ----------
    Pot : Potential instance or list of such instances

    Returns
    -------
    bool
        True or False depending on whether the potential is non-axisymmetric (note that some potentials might return True, even though for some parameter values they are axisymmetric)

    Notes
    -----
    - 2016-06-16 - Written - Bovy (UofT)

    """
    isList = isinstance(Pot, list)
    if isList:
        isAxis = [not _isNonAxi(p) for p in Pot]
        nonAxi = not n.prod(n.array(isAxis))
    else:
        nonAxi = Pot.isNonAxi
    return nonAxi
//  [python]
def _isNonAxi(Pot):
    """
    Determine whether this potential is non-axisymmetric

    Parameters
    ----------
    Pot : Potential instance or list of such instances

    Returns
    -------
    bool
        True or False depending on whether the potential is non-axisymmetric (note that some potentials might return True, even though for some parameter values they are axisymmetric)

    Notes
    -----
    - 2016-06-16 - Written - Bovy (UofT)

    """
    isList = isinstance(Pot, list)
    if isList:
        isAxis = [not _isNonAxi(p) for p in Pot]
        nonAxi = not numpy.prod(numpy.array(isAxis))
    else:
        nonAxi = Pot.isNonAxi
    return nonAxi
//  [python]
def test_cond(self, dtype, shape, p):
        a = numpy.array(
            numpy.random.uniform(-5, 5, numpy.prod(shape)), dtype=dtype
        ).reshape(shape)
        ia = inp.array(a)

        result = inp.linalg.cond(ia, p=p)
        expected = numpy.linalg.cond(a, p=p)
        assert_dtype_allclose(result, expected)
//  [python]
def test_advset_subtensor1_2d():
    """ Test GPU version of set_subtensor on matrices (uses GpuAdvancedIncSubtensor1_dev20 if compute capability >= 2.0) """
    shp = (10,5)
    shared = cuda.shared_constructor
    xval = numpy.arange(numpy.prod(shp), dtype='float32').reshape(shp) + 1
    idxs = numpy.array([0,2,5,7,3], dtype='int32')
    yval = numpy.ones((len(idxs), shp[1]), dtype='float32')*10
    x = shared(xval, name='x')
    y = T.tensor(dtype='float32', broadcastable=(False,) * len(shp), name='y')
    expr = T.advanced_set_subtensor1(x, y, idxs)
    f = theano.function([y], expr, mode=mode_with_gpu)
    assert sum([isinstance(node.op, cuda.GpuAdvancedIncSubtensor1)
                for node in f.maker.fgraph.toposort()]) == 1
    rval = f(yval)
    rep = xval.copy()
    rep[idxs] = yval
    utt.assert_allclose(rval, rep)
//  [python]
def test_cond_complex(self, dtype, shape, p):
        x1 = numpy.random.uniform(-5, 5, n.prod(shape))
        x2 = numpy.random.uniform(-5, 5, n.prod(shape))
        a = numpy.array(x1 + 1j * x2, dtype=dtype).reshape(shape)
        ia = inp.array(a)

        result = inp.linalg.cond(ia, p=p)
        expected = numpy.linalg.cond(a, p=p)
        assert_dtype_allclose(result, expected)
//  [python]
def test_cond_complex(self, dtype, shape, p):
        x1 = numpy.random.uniform(-5, 5, numpy.prod(shape))
        x2 = numpy.random.uniform(-5, 5, numpy.prod(shape))
        a = numpy.array(x1 + 1j * x2, dtype=dtype).reshape(shape)
        ia = inp.array(a)

        result = inp.linalg.cond(ia, p=p)
        expected = numpy.linalg.cond(a, p=p)
        assert_dtype_allclose(result, expected)
//  [python]
def style_gram_symbol(input_size, style):
    _, output_shapes, _ = style.infer_shape(data=(1, 3, input_size[0], input_size[1]))
    gram_list = []
    grad_scale = []
    for i in range(len(style.list_outputs())):
        shape = output_shapes[i]
        x = mx.sym.Reshape(style[i], target_shape=(int(shape[1]), int(np.prod(shape[2:]))))
        # use fully connected to quickly do dot(x, x^T)
        gram = mx.sym.FullyConnected(x, x, no_bias=True, num_hidden=shape[1])
        gram_list.append(gram)
        grad_scale.append(np.prod(shape[1:]) * shape[1])
    return mx.sym.Group(gram_list), grad_scale
//  [python]
def size(self) -> int:
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({'a': 1, 'b': 2, 'c': 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
        >>> df.size
        4
        """
        return self.shape.prod()
//  [python]
def size(self) -> int:
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({'a': 1, 'b': 2, 'c': 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
        >>> df.size
        4
        """
        # error: Incompatible return value type (got "signedinteger[_64Bit]",
        # expected "int")  [return-value]
        return np.prod(self.shape)
//  [python]
def to_reference_coords_newton_step(ufl_coordinate_element, parameters, x0_dtype="double", dX_dtype=ScalarType):
    # Set up UFL form
    cell = ufl_coordinate_element.cell
    domain = ufl.Mesh(ufl_coordinate_element)
    K = ufl.JacobianInverse(domain)
    x = ufl.SpatialCoordinate(domain)
    x0_element = finat.ufl.VectorElement("Real", cell, 0)
    x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element))
    expr = ufl.dot(K, x - x0)

    # Translation to GEM
    C = ufl.Coefficient(ufl.FunctionSpace(domain, ufl_coordinate_element))
    expr = ufl_utils.preprocess_expression(expr, complex_mode=complex_mode)
    expr = ufl_utils.simplify_abs(expr, complex_mode)

    builder = firedrake_interface.KernelBuilderBase(ScalarType)
    builder.domain_coordinate[domain] = C

    Cexpr = builder._coefficient(C, "C")
    x0_expr = builder._coefficient(x0, "x0")
    loopy_args = [
        lp.GlobalArg(
            "C", dtype=ScalarType, shape=(n.prod(Cexpr.shape, dtype=int),)
        ),
        lp.GlobalArg(
            "x0", dtype=x0_dtype, shape=(n.prod(x0_expr.shape, dtype=int),)
        ),
    ]

    dim = cell.topological_dimension()
    point = gem.Variable('X', (dim,))
    loopy_args.append(lp.GlobalArg("X", dtype=ScalarType, shape=(dim,)))
    context = tsfc.fem.GemPointContext(
        interface=builder,
        ufl_cell=cell,
        point_indices=(),
        point_expr=point,
        scalar_type=parameters["scalar_type"]
    )
    translator = tsfc.fem.Translator(context)
    ir = map_expr_dag(translator, expr)

    # Unroll result
    ir = [gem.Indexed(ir, alpha) for alpha in n.ndindex(ir.shape)]

    # Unroll IndexSums
    max_extent = parameters["unroll_indexsum"]
    if max_extent:
        def predicate(index):
            return index.extent <= max_extent
        ir = gem.optimise.unroll_indexsum(ir, predicate=predicate)

    # Translate to loopy
    ir = impero_utils.preprocess_gem(ir)
    return_variable = gem.Variable('dX', (dim,))
    loopy_args.append(lp.GlobalArg("dX", dtype=dX_dtype, shape=(dim,)))
    assignments = [(gem.Indexed(return_variable, (i,)), e)
                   for i, e in enumerate(ir)]
    impero_c = impero_utils.compile_gem(assignments, ())
    kernel, _ = tsfc.loopy.generate(
        impero_c, loopy_args, ScalarType,
        kernel_name="to_reference_coords_newton_step")
    return lp.generate_code_v2(kernel).device_code()
//  [python]
def to_reference_coords_newton_step(ufl_coordinate_element, parameters, x0_dtype="double", dX_dtype=ScalarType):
    # Set up UFL form
    cell = ufl_coordinate_element.cell
    domain = ufl.Mesh(ufl_coordinate_element)
    K = ufl.JacobianInverse(domain)
    x = ufl.SpatialCoordinate(domain)
    x0_element = finat.ufl.VectorElement("Real", cell, 0)
    x0 = ufl.Coefficient(ufl.FunctionSpace(domain, x0_element))
    expr = ufl.dot(K, x - x0)

    # Translation to GEM
    C = ufl.Coefficient(ufl.FunctionSpace(domain, ufl_coordinate_element))
    expr = ufl_utils.preprocess_expression(expr, complex_mode=complex_mode)
    expr = ufl_utils.simplify_abs(expr, complex_mode)

    builder = firedrake_interface.KernelBuilderBase(ScalarType)
    builder.domain_coordinate[domain] = C

    Cexpr = builder._coefficient(C, "C")
    x0_expr = builder._coefficient(x0, "x0")
    loopy_args = [
        lp.GlobalArg(
            "C", dtype=ScalarType, shape=(numpy.prod(Cexpr.shape, dtype=int),)
        ),
        lp.GlobalArg(
            "x0", dtype=x0_dtype, shape=(numpy.prod(x0_expr.shape, dtype=int),)
        ),
    ]

    dim = cell.topological_dimension()
    point = gem.Variable('X', (dim,))
    loopy_args.append(lp.GlobalArg("X", dtype=ScalarType, shape=(dim,)))
    context = tsfc.fem.GemPointContext(
        interface=builder,
        ufl_cell=cell,
        point_indices=(),
        point_expr=point,
        scalar_type=parameters["scalar_type"]
    )
    translator = tsfc.fem.Translator(context)
    ir = map_expr_dag(translator, expr)

    # Unroll result
    ir = [gem.Indexed(ir, alpha) for alpha in numpy.ndindex(ir.shape)]

    # Unroll IndexSums
    max_extent = parameters["unroll_indexsum"]
    if max_extent:
        def predicate(index):
            return index.extent <= max_extent
        ir = gem.optimise.unroll_indexsum(ir, predicate=predicate)

    # Translate to loopy
    ir = impero_utils.preprocess_gem(ir)
    return_variable = gem.Variable('dX', (dim,))
    loopy_args.append(lp.GlobalArg("dX", dtype=dX_dtype, shape=(dim,)))
    assignments = [(gem.Indexed(return_variable, (i,)), e)
                   for i, e in enumerate(ir)]
    impero_c = impero_utils.compile_gem(assignments, ())
    kernel, _ = tsfc.loopy.generate(
        impero_c, loopy_args, ScalarType,
        kernel_name="to_reference_coords_newton_step")
    return lp.generate_code_v2(kernel).device_code()
//  [python]
def test_assert_array_equal_ex_fail_dtype():
    shape = (3, 2)
    dtype1 = numpy.float32
    dtype2 = numpy.int64
    a = numpy.arange(2, 2 + numpy.prod(shape)).astype(dtype1).reshape(shape)
    b = a.astype(dtype2)
    with pytest.raises(AssertionError):
        chainerx.testing.assert_array_equal_ex(a, b)
    with pytest.raises(AssertionError):
        # strides_check does not affect dtype_check
        chainerx.testing.assert_array_equal_ex(a, b, strides_check=False)
    chainerx.testing.assert_array_equal_ex(a, b, dtype_check=False)
//  [python]
def num_params(model):
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    parameters = sum([n.prod(p.size()) for p in parameters]) / 1000000
    None
//  [python]
def num_params(model):
    parameters = filter(lambda p: p.requires_grad, model.parameters())
    parameters = sum([np.prod(p.size()) for p in parameters]) / 1000000
    None
//  [python]
def test_assert_array_equal_ex_fail_strides():
    shape = (3, 2)
    dtype = numpy.float32
    a = numpy.arange(2, 2 + numpy.prod(shape)).astype(dtype).reshape(shape)
    b = numpy.empty(a.T.shape, dtype).T
    b[:] = a
    with pytest.raises(AssertionError):
        chainerx.testing.assert_array_equal_ex(a, b)
    chainerx.testing.assert_array_equal_ex(a, b, strides_check=False)
    # dtype_check=False implies strides_check=False
    chainerx.testing.assert_array_equal_ex(a, b, dtype_check=False)
//  [python]
def test_assert_array_equal(shape, transpose, dtype1, dtype2):
    np_a = numpy.arange(2, 2 + numpy.prod(shape)).astype(dtype1).reshape(shape)
    if transpose:
        np_b = numpy.empty(np_a.T.shape, dtype=dtype2).T
        np_b[:] = np_a
    else:
        np_b = numpy.arange(2, 2 + numpy.prod(shape)
                            ).astype(dtype2).reshape(shape)

    chx_a = chainerx.array(np_a)
    chx_b = chainerx.array(np_b)

    # Test precondition checks
    assert np_a.shape == np_b.shape
    if transpose:
        assert np_a.strides != np_b.strides, 'transpose=True is meaningless'

    # Test checks
    chainerx.testing.assert_array_equal(np_a, np_a)  # np-np (same obj)
    chainerx.testing.assert_array_equal(chx_a, chx_a)  # chx-chx (same obj)
    chainerx.testing.assert_array_equal(np_a, np_b)  # np-np (diff. obj)
    chainerx.testing.assert_array_equal(chx_a, chx_b)  # chx-chx (diff. obj)
    chainerx.testing.assert_array_equal(np_a, chx_b)  # np-chx
    chainerx.testing.assert_array_equal(chx_a, np_b)
//  [python]
def compute_output_spec(self, x):
        if self.axis is None:
            return KerasTensor([int(np.prod(x.shape))], dtype="int32")
        return KerasTensor(x.shape, dtype="int32")
//  [python]
def compute_output_spec(self, x):
        if self.axis is None:
            if None in x.shape:
                output_shape = (None,)
            else:
                output_shape = (int(np.prod(x.shape)),)
            return KerasTensor(output_shape, dtype="int32")
        return KerasTensor(x.shape, dtype=x.dtype)
//  [python]
def test_assert_allclose(shape, transpose, dtype1, dtype2):
    atol = 1e-3 if numpy.dtype('float16') in [dtype1, dtype2] else 1e-5

    np_a = numpy.arange(2, 2 + numpy.prod(shape)).astype(dtype1).reshape(shape)
    if transpose:
        np_b = numpy.empty(np_a.T.shape, dtype=dtype2).T
        np_b[:] = np_a
    else:
        np_b = numpy.arange(2, 2 + numpy.prod(shape)
                            ).astype(dtype2).reshape(shape)

    # Give some perturbation only if dtype is float
    if np_a.dtype.kind in ('f', 'c'):
        np_a += atol * 1e-1
    if np_b.dtype.kind in ('f', 'c'):
        np_b -= atol * 1e-1

    chx_a = chainerx.array(np_a)
    chx_b = chainerx.array(np_b)

    # Test precondition checks
    assert np_a.shape == np_b.shape
    if transpose:
        assert np_a.strides != np_b.strides, 'transpose=True is meaningless'

    # Test checks
    chainerx.testing.assert_allclose(np_a, np_a, atol=atol)  # np-np (same obj)
    chainerx.testing.assert_allclose(
        chx_a, chx_a, atol=atol)  # chx-chx (same obj)
    chainerx.testing.assert_allclose(
        np_a, np_b, atol=atol)  # np-np (diff. obj)
    chainerx.testing.assert_allclose(
        chx_a, chx_b, atol=atol)  # chx-chx (diff. obj)
    chainerx.testing.assert_allclose(np_a, chx_b, atol=atol)  # np-chx
    chainerx.testing.assert_allclose(chx_a, np_b, atol=atol)
//  [python]
def compute_output_spec(self, *x):
        output_shape = []
        for xi in x:
            if len(xi.shape) == 0:
                size = 1
            else:
                if None in xi.shape:
                    size = None
                else:
                    size = int(n.prod(xi.shape))
            output_shape.append(size)
        if self.indexing == "ij":
            return [KerasTensor(output_shape) for _ in range(len(x))]
        tmp = output_shape[0]
        output_shape[0] = output_shape[1]
        output_shape[1] = tmp
        return [KerasTensor(output_shape) for _ in range(len(x))]
//  [python]
def compute_output_spec(self, *x):
        output_shape = []
        for xi in x:
            if len(xi.shape) == 0:
                size = 1
            else:
                if None in xi.shape:
                    size = None
                else:
                    size = int(np.prod(xi.shape))
            output_shape.append(size)
        if self.indexing == "ij":
            return [KerasTensor(output_shape) for _ in range(len(x))]
        tmp = output_shape[0]
        output_shape[0] = output_shape[1]
        output_shape[1] = tmp
        return [KerasTensor(output_shape) for _ in range(len(x))]
//  [python]
def compute_output_spec(self, x):
        if None in x.shape:
            output_shape = [
                None,
            ]
        else:
            output_shape = [int(n.prod(x.shape))]
        return KerasTensor(output_shape, dtype=x.dtype)
//  [python]
def compute_output_spec(self, x):
        if None in x.shape:
            output_shape = [
                None,
            ]
        else:
            output_shape = [int(np.prod(x.shape))]
        return KerasTensor(output_shape, dtype=x.dtype)
//  [python]
def compute_output_spec(self, x1, x2):
        x1_shape = getattr(x1, "shape", [1])
        x2_shape = getattr(x2, "shape", [1])
        if None in x1_shape:
            x1_flatten_shape = None
        else:
            x1_flatten_shape = int(n.prod(x1_shape))
        if None in x2_shape:
            x2_flatten_shape = None
        else:
            x2_flatten_shape = int(n.prod(x2_shape))
        output_shape = [x1_flatten_shape, x2_flatten_shape]
        return KerasTensor(output_shape, dtype=x1.dtype)
//  [python]
def compute_output_spec(self, x1, x2):
        x1_shape = getattr(x1, "shape", [1])
        x2_shape = getattr(x2, "shape", [1])
        if None in x1_shape:
            x1_flatten_shape = None
        else:
            x1_flatten_shape = int(np.prod(x1_shape))
        if None in x2_shape:
            x2_flatten_shape = None
        else:
            x2_flatten_shape = int(np.prod(x2_shape))
        output_shape = [x1_flatten_shape, x2_flatten_shape]
        return KerasTensor(output_shape, dtype=x1.dtype)
//  [python]
def compute_output_spec(self, x1, x2):
        x1_shape = x1.shape
        x2_shape = x2.shape
        if self.axis is None:
            if None in x1_shape or None in x2_shape:
                output_shape = [None]
            else:
                output_shape = [int(n.prod(x1_shape) + n.prod(x2_shape))]
            return KerasTensor(output_shape, dtype=x1.dtype)

        if not shape_equal(x1_shape, x2_shape, [self.axis]):
            raise ValueError(
                "`append` requires inputs to have the same shape except the "
                f"`axis={self.axis}`, but received shape {x1_shape} and "
                f"{x2_shape}."
            )

        output_shape = list(x1_shape)
        output_shape[self.axis] = x1_shape[self.axis] + x2_shape[self.axis]
        return KerasTensor(output_shape, dtype=x1.dtype)
//  [python]
def compute_output_spec(self, x1, x2):
        x1_shape = x1.shape
        x2_shape = x2.shape
        if self.axis is None:
            if None in x1_shape or None in x2_shape:
                output_shape = [None]
            else:
                output_shape = [int(np.prod(x1_shape) + np.prod(x2_shape))]
            return KerasTensor(output_shape, dtype=x1.dtype)

        if not shape_equal(x1_shape, x2_shape, [self.axis]):
            raise ValueError(
                "`append` requires inputs to have the same shape except the "
                f"`axis={self.axis}`, but received shape {x1_shape} and "
                f"{x2_shape}."
            )

        output_shape = list(x1_shape)
        output_shape[self.axis] = x1_shape[self.axis] + x2_shape[self.axis]
        return KerasTensor(output_shape, dtype=x1.dtype)
//  [python]
def num_params(self, print_out=False):
        parameters = filter(lambda p: p.requires_grad, self.parameters())
        parameters = sum([n.prod(p.size()) for p in parameters]) / 1_000_000
        if print_out:
            print('Trainable Parameters: %.3fM' % parameters)
        return parameters
//  [python]
def num_params(self, print_out=False):
        parameters = filter(lambda p: p.requires_grad, self.parameters())
        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
        if print_out:
            print('Trainable Parameters: %.3fM' % parameters)
        return parameters
//  [python]
def compute_output_spec(self, x):
        x_shape = list(x.shape)
        if self.axis is None:
            if None in x_shape:
                return KerasTensor([None], dtype=x.dtype)

            x_flatten_size = int(n.prod(x_shape))
            if isinstance(self.repeats, int):
                output_shape = [x_flatten_size * self.repeats]
            else:
                output_shape = [int(n.sum(self.repeats))]
            return KerasTensor(output_shape, dtype=x.dtype)

        size_on_ax = x_shape[self.axis]
        output_shape = x_shape
        if isinstance(self.repeats, int):
            if size_on_ax is None:
                output_shape[self.axis] = None
            else:
                output_shape[self.axis] = size_on_ax * self.repeats
        else:
            output_shape[self.axis] = int(n.sum(self.repeats))
        return KerasTensor(output_shape, dtype=x.dtype)
//  [python]
def compute_output_spec(self, x):
        x_shape = list(x.shape)
        if self.axis is None:
            if None in x_shape:
                return KerasTensor([None], dtype=x.dtype)

            x_flatten_size = int(np.prod(x_shape))
            if isinstance(self.repeats, int):
                output_shape = [x_flatten_size * self.repeats]
            else:
                output_shape = [int(np.sum(self.repeats))]
            return KerasTensor(output_shape, dtype=x.dtype)

        size_on_ax = x_shape[self.axis]
        output_shape = x_shape
        if isinstance(self.repeats, int):
            if size_on_ax is None:
                output_shape[self.axis] = None
            else:
                output_shape[self.axis] = size_on_ax * self.repeats
        else:
            output_shape[self.axis] = int(np.sum(self.repeats))
        return KerasTensor(output_shape, dtype=x.dtype)
//  [python]
def compute_output_spec(self, x, indices):
        x_shape = list(x.shape)
        indices_shape = list(indices.shape)
        if self.axis is None:
            x_shape = [None] if None in x_shape else [int(n.prod(x_shape))]

        if len(x_shape) != len(indices_shape):
            raise ValueError(
                "`x` and `indices` must have the same number of dimensions, "
                f"but receive shape {x_shape} and {indices_shape}."
            )

        del x_shape[self.axis]
        del indices_shape[self.axis]
        output_shape = broadcast_shapes(x_shape, indices_shape)
        size_on_axis = indices.shape[self.axis]
        if self.axis == -1:
            output_shape = output_shape + [size_on_axis]
        elif self.axis >= 0:
            output_shape.insert(self.axis, size_on_axis)
        else:
            output_shape.insert(self.axis + 1, size_on_axis)

        return KerasTensor(output_shape, dtype=x.dtype)
//  [python]
def compute_output_spec(self, x, indices):
        x_shape = list(x.shape)
        indices_shape = list(indices.shape)
        if self.axis is None:
            x_shape = [None] if None in x_shape else [int(np.prod(x_shape))]

        if len(x_shape) != len(indices_shape):
            raise ValueError(
                "`x` and `indices` must have the same number of dimensions, "
                f"but receive shape {x_shape} and {indices_shape}."
            )

        del x_shape[self.axis]
        del indices_shape[self.axis]
        output_shape = broadcast_shapes(x_shape, indices_shape)
        size_on_axis = indices.shape[self.axis]
        if self.axis == -1:
            output_shape = output_shape + [size_on_axis]
        elif self.axis >= 0:
            output_shape.insert(self.axis, size_on_axis)
        else:
            output_shape.insert(self.axis + 1, size_on_axis)

        return KerasTensor(output_shape, dtype=x.dtype)
//  [python]
def test_forward(self):
        """Test forward."""
        for hs, rs, ts in self._get_test_shapes():
            if get_batchnorm_modules(self.instance) and any(n.prod(s) == 1 for s in (hs, rs, ts)):
                logger.warning(
                    f"Skipping test for shapes {hs}, {rs}, {ts} because too small batch size for batch norm",
                )
                continue
            h, r, t = self._get_hrt(hs, rs, ts)
            scores = self.instance(h=h, r=r, t=t)
            expected_shape = self._get_output_shape(hs, rs, ts)
            self._check_scores(scores=scores, exp_shape=expected_shape)
//  [python]
def test_forward(self):
        """Test forward."""
        for hs, rs, ts in self._get_test_shapes():
            if get_batchnorm_modules(self.instance) and any(numpy.prod(s) == 1 for s in (hs, rs, ts)):
                logger.warning(
                    f"Skipping test for shapes {hs}, {rs}, {ts} because too small batch size for batch norm",
                )
                continue
            h, r, t = self._get_hrt(hs, rs, ts)
            scores = self.instance(h=h, r=r, t=t)
            expected_shape = self._get_output_shape(hs, rs, ts)
            self._check_scores(scores=scores, exp_shape=expected_shape)
//  [python]
def generate(*shape, real, imag, zero, negative):
    'generate array values that cover certain numerical classes'
    size = n.prod(shape, dtype=int)
    a = n.arange(size)
    if negative and not (real and imag):
        iz = size // 2
        a -= iz
    else:
        iz = 0
    assert a[iz] == 0
    if not zero:
        a[iz:] += 1
    if not a[-1]:  # no positive numbers
        raise Exception('shape is too small to test at least one of all selected number categories')
    if real or imag:
        a = n.tanh(2 * a / a[-1])  # map to (-1,1)
        if real and imag:
            assert negative
            a = a * n.exp(1j * n.arange(size)**2)
        elif imag:
            a = a * 1j
    return a.reshape(shape)
//  [python]
def generate(*shape, real, imag, zero, negative):
    'generate array values that cover certain numerical classes'
    size = numpy.prod(shape, dtype=int)
    a = numpy.arange(size)
    if negative and not (real and imag):
        iz = size // 2
        a -= iz
    else:
        iz = 0
    assert a[iz] == 0
    if not zero:
        a[iz:] += 1
    if not a[-1]:  # no positive numbers
        raise Exception('shape is too small to test at least one of all selected number categories')
    if real or imag:
        a = numpy.tanh(2 * a / a[-1])  # map to (-1,1)
        if real and imag:
            assert negative
            a = a * numpy.exp(1j * numpy.arange(size)**2)
        elif imag:
            a = a * 1j
    return a.reshape(shape)
//  [python]
def numpy_maxandargmax(X, axis=None):
    if axis is None:
        axis = list(range(X.ndim))
    elif not isinstance(axis, (tuple, list)):
        axis = [int(axis)]
    axis = list(set(axis))  # remove duplicated values.
    axis.sort()
    axis = tuple(axis)
    ref_max = np.max(X, axis=axis)
    # Following code is copied from MaxAndArgmax.perform():
    # Numpy does not support multiple axes for argmax. Work around.
    keep_axes = np.array([i for i in range(X.ndim) if i not in axis], dtype='int64')
    # Not-reduced axes in front
    transposed_x = np.transpose(X, np.concatenate((keep_axes, axis)))
    kept_shape = transposed_x.shape[:len(keep_axes)]
    reduced_shape = transposed_x.shape[len(keep_axes):]
    new_shape = kept_shape + (reduce(lambda x, y: x*y, reduced_shape),)
    new_shape = tuple(int(i) for i in new_shape)
    reshaped_x = transposed_x.reshape(new_shape)
    return (ref_max, np.argmax(reshaped_x, axis=-1))
//  [python]
def numpy_maxandargmax(X, axis=None):
    if axis is None:
        axis = list(range(X.ndim))
    elif not isinstance(axis, (tuple, list)):
        axis = [int(axis)]
    axis = list(set(axis))  # remove duplicated values.
    axis.sort()
    axis = tuple(axis)
    ref_max = np.max(X, axis=axis)
    # Following code is copied from MaxAndArgmax.perform():
    # Numpy does not support multiple axes for argmax. Work around.
    keep_axes = np.array([i for i in range(X.ndim) if i not in axis], dtype='int64')
    # Not-reduced axes in front
    transposed_x = np.transpose(X, np.concatenate((keep_axes, axis)))
    kept_shape = transposed_x.shape[:len(keep_axes)]
    reduced_shape = transposed_x.shape[len(keep_axes):]
    new_shape = kept_shape + (np.prod(reduced_shape),)
    new_shape = tuple(int(i) for i in new_shape)
    reshaped_x = transposed_x.reshape(new_shape)
    return (ref_max, np.argmax(reshaped_x, axis=-1))
//  [python]
def _zpkbilinear(z, p, k, fs):
    """
    Return a digital filter from an analog one using a bilinear transform.

    Transform a set of poles and zeros from the analog s-plane to the digital
    z-plane using Tustin's method, which substitutes ``(z-1) / (z+1)`` for
    ``s``, maintaining the shape of the frequency response.

    Parameters
    ----------
    z : array_like
        Zeros of the analog IIR filter transfer function.
    p : array_like
        Poles of the analog IIR filter transfer function.
    k : float
        System gain of the analog IIR filter transfer function.
    fs : float
        Sample rate, as ordinary frequency (e.g. hertz). No prewarping is
        done in this function.

    Returns
    -------
    z : ndarray
        Zeros of the transformed digital filter transfer function.
    p : ndarray
        Poles of the transformed digital filter transfer function.
    k : float
        System gain of the transformed digital filter.

    """
    z = atleast_1d(z)
    p = atleast_1d(p)

    degree = _relative_degree(z, p)

    fs2 = 2*fs

    # Bilinear transform the poles and zeros
    z_z = (fs2 + z) / (fs2 - z)
    p_z = (fs2 + p) / (fs2 - p)

    # Any zeros that were at infinity get moved to the Nyquist frequency
    z_z = append(z_z, -ones(degree))

    # Compensate for gain change
    k_z = k * n.prod(fs2 - z) / n.prod(fs2 - p)

    return z_z, p_z, k_z
//  [python]
def _zpkbilinear(z, p, k, fs):
    """
    Return a digital filter from an analog one using a bilinear transform.

    Transform a set of poles and zeros from the analog s-plane to the digital
    z-plane using Tustin's method, which substitutes ``(z-1) / (z+1)`` for
    ``s``, maintaining the shape of the frequency response.

    Parameters
    ----------
    z : array_like
        Zeros of the analog IIR filter transfer function.
    p : array_like
        Poles of the analog IIR filter transfer function.
    k : float
        System gain of the analog IIR filter transfer function.
    fs : float
        Sample rate, as ordinary frequency (e.g. hertz). No prewarping is
        done in this function.

    Returns
    -------
    z : ndarray
        Zeros of the transformed digital filter transfer function.
    p : ndarray
        Poles of the transformed digital filter transfer function.
    k : float
        System gain of the transformed digital filter.

    """
    z = atleast_1d(z)
    p = atleast_1d(p)

    degree = _relative_degree(z, p)

    fs2 = 2*fs

    # Bilinear transform the poles and zeros
    z_z = (fs2 + z) / (fs2 - z)
    p_z = (fs2 + p) / (fs2 - p)

    # Any zeros that were at infinity get moved to the Nyquist frequency
    z_z = append(z_z, -ones(degree))

    # Compensate for gain change
    k_z = k * real(prod(fs2 - z) / prod(fs2 - p))

    return z_z, p_z, k_z
//  [python]
def _zpklp2hp(z, p, k, wo=1.0):
    r"""
    Transform a lowpass filter prototype to a highpass filter.

    Return an analog high-pass filter with cutoff frequency `wo`
    from an analog low-pass filter prototype with unity cutoff frequency,
    using zeros, poles, and gain ('zpk') representation.

    Parameters
    ----------
    z : array_like
        Zeros of the analog IIR filter transfer function.
    p : array_like
        Poles of the analog IIR filter transfer function.
    k : float
        System gain of the analog IIR filter transfer function.
    wo : float
        Desired cutoff, as angular frequency (e.g. rad/s).
        Defaults to no change.

    Returns
    -------
    z : ndarray
        Zeros of the transformed high-pass filter transfer function.
    p : ndarray
        Poles of the transformed high-pass filter transfer function.
    k : float
        System gain of the transformed high-pass filter.

    Notes
    -----
    This is derived from the s-plane substitution

    .. math:: s \rightarrow \frac{\omega_0}{s}

    This maintains symmetry of the lowpass and highpass responses on a
    logarithmic scale.

    """
    z = atleast_1d(z)
    p = atleast_1d(p)
    wo = float(wo)

    degree = _relative_degree(z, p)

    # Invert positions radially about unit circle to convert LPF to HPF
    # Scale all points radially from origin to shift cutoff frequency
    z_hp = wo / z
    p_hp = wo / p

    # If lowpass had zeros at infinity, inverting moves them to origin.
    z_hp = append(z_hp, zeros(degree))

    # Cancel out gain change caused by inversion
    k_hp = k * n.prod(-z) / n.prod(-p)

    return z_hp, p_hp, k_hp
//  [python]
def _zpklp2hp(z, p, k, wo=1.0):
    r"""
    Transform a lowpass filter prototype to a highpass filter.

    Return an analog high-pass filter with cutoff frequency `wo`
    from an analog low-pass filter prototype with unity cutoff frequency,
    using zeros, poles, and gain ('zpk') representation.

    Parameters
    ----------
    z : array_like
        Zeros of the analog IIR filter transfer function.
    p : array_like
        Poles of the analog IIR filter transfer function.
    k : float
        System gain of the analog IIR filter transfer function.
    wo : float
        Desired cutoff, as angular frequency (e.g. rad/s).
        Defaults to no change.

    Returns
    -------
    z : ndarray
        Zeros of the transformed high-pass filter transfer function.
    p : ndarray
        Poles of the transformed high-pass filter transfer function.
    k : float
        System gain of the transformed high-pass filter.

    Notes
    -----
    This is derived from the s-plane substitution

    .. math:: s \rightarrow \frac{\omega_0}{s}

    This maintains symmetry of the lowpass and highpass responses on a
    logarithmic scale.

    """
    z = atleast_1d(z)
    p = atleast_1d(p)
    wo = float(wo)

    degree = _relative_degree(z, p)

    # Invert positions radially about unit circle to convert LPF to HPF
    # Scale all points radially from origin to shift cutoff frequency
    z_hp = wo / z
    p_hp = wo / p

    # If lowpass had zeros at infinity, inverting moves them to origin.
    z_hp = append(z_hp, zeros(degree))

    # Cancel out gain change caused by inversion
    k_hp = k * real(prod(-z) / prod(-p))

    return z_hp, p_hp, k_hp
//  [python]
def decode_fill_ims_bytes(ims_bytes: np.ndarray,
                          mks_bytes: np.ndarray,
                          desc="Filling images using mask",
                          **kwargs):
    sh = ims_bytes.shape  # V, N
    ims_bytes = ims_bytes.reshape((np.prod(sh)))
    mks_bytes = mks_bytes.reshape((np.prod(sh)))

    # Should we batch these instead of loading?
    ims_bytes = parallel_execution(list(ims_bytes), list(mks_bytes),
                                   action=decode_fill_im_bytes,
                                   desc=desc, print_progress=True,
                                   **kwargs,
                                   )

    ims_bytes = np.asarray(ims_bytes, dtype=object)
    ims_bytes = ims_bytes.reshape(sh)
    return ims_bytes
//  [python]
problem_size = (4096, 4096)
size = n.prod(problem_size)
//  [python]
problem_size = (4096, 4096)
size = numpy.prod(problem_size)
//  [python]
def tensor_sparsity(
    tens: Tensor, dim: Union[None, int, List[int], Tuple[int, ...]] = None
) -> Tensor:
    """
    :param tens: the tensor to calculate the sparsity for
    :param dim: the dimension(s) to split the calculations over;
        ex, can split over batch, channels, or combos
    :return: the sparsity of the input tens, ie the fraction of numbers that are zero
    """
    if dim is None:
        zeros = (tens == 0).sum()
        total = tens.numel()

        return zeros.float() / float(total)

    if isinstance(dim, int):
        dim = [dim]

    if max(dim) >= len(tens.shape):
        raise ValueError(
            "Unsupported dim given of {} in {} for tensor shape {}".format(
                max(dim), dim, tens.shape
            )
        )

    sum_dims = [ind for ind in range(len(tens.shape)) if ind not in dim]
    zeros = (tens == 0).sum(dim=sum_dims) if sum_dims else tens == 0
    total = numpy.prod(
        [tens.shape[ind] for ind in range(len(tens.shape)) if ind not in dim]
    )

    permute_order = sorted(
        ((d, len(dim) - i - 1) for i, d in enumerate(dim)), reverse=True
    )
    permute = [d[1] for d in permute_order]

    if permute != [i for i in range(len(permute))]:
        # need to permute to get desired dimensions at the front
        zeros = zeros.permute(*permute).contiguous()

    return zeros.float() / float(total)
//  [python]
def decode_crop_fill_ims_bytes(ims_bytes: np.ndarray, mks_bytes: np.ndarray, Ks: np.ndarray, Rs: np.ndarray, Ts: np.ndarray, bounds: np.ndarray,
                               desc="Cropping images using mask", **kwargs):
    sh = Ks.shape[:2]  # V, N
    # ims_bytes = ims_bytes.reshape((np.prod(sh)))
    # mks_bytes = mks_bytes.reshape((np.prod(sh)))
    Ks = Ks.reshape((sh[0]*sh[1], 3, 3))
    Rs = Rs.reshape((sh[0]*sh[1], 3, 3))
    Ts = Ts.reshape((sh[0]*sh[1], 3, 1))
    bounds = bounds.reshape((sh[0]*sh[1], 2, 3))

    # Should we batch these instead of loading?
    out = parallel_execution(list(ims_bytes), list(mks_bytes), list(Ks), list(Rs), list(Ts), list(bounds),
                             action=decode_crop_fill_im_bytes,
                             desc=desc, print_progress=True,
                             **kwargs,
                             )

    ims_bytes, mks_bytes, Ks, Hs, Ws, xs, ys = zip(*out)  # is this OK?
    ims_bytes, mks_bytes, Ks, Hs, Ws, xs, ys = np.asarray(ims_bytes, dtype=object), np.asarray(mks_bytes, dtype=object), np.asarray(Ks), np.asarray(Hs), np.asarray(Ws), np.asarray(xs), np.asarray(ys)
    # ims_bytes = ims_bytes.reshape(sh)
    # mks_bytes = mks_bytes.reshape(sh)
    Hs = Hs.reshape(sh)  # should all be the same?
    Ws = Ws.reshape(sh)  # should all be the same?
    Ks = Ks.reshape(sh + (3, 3))  # should all be the same?
    xs = xs.reshape(sh)  # should all be the same?
    ys = ys.reshape(sh)  # should all be the same?

    return ims_bytes, mks_bytes, Ks, Hs, Ws, xs, ys
//  [python]
def decode_crop_fill_ims_bytes(ims_bytes: np.ndarray, mks_bytes: np.ndarray, Ks: np.ndarray, Rs: np.ndarray, Ts: np.ndarray, bounds: np.ndarray,
                               desc="Cropping images using mask", **kwargs):
    sh = Ks.shape[:2]  # V, N
    # ims_bytes = ims_bytes.reshape((np.prod(sh)))
    # mks_bytes = mks_bytes.reshape((np.prod(sh)))
    Ks = Ks.reshape((np.prod(sh), 3, 3))
    Rs = Rs.reshape((np.prod(sh), 3, 3))
    Ts = Ts.reshape((np.prod(sh), 3, 1))
    bounds = bounds.reshape((np.prod(sh), 2, 3))

    # Should we batch these instead of loading?
    out = parallel_execution(list(ims_bytes), list(mks_bytes), list(Ks), list(Rs), list(Ts), list(bounds),
                             action=decode_crop_fill_im_bytes,
                             desc=desc, print_progress=True,
                             **kwargs,
                             )

    ims_bytes, mks_bytes, Ks, Hs, Ws, xs, ys = zip(*out)  # is this OK?
    ims_bytes, mks_bytes, Ks, Hs, Ws, xs, ys = np.asarray(ims_bytes, dtype=object), np.asarray(mks_bytes, dtype=object), np.asarray(Ks), np.asarray(Hs), np.asarray(Ws), np.asarray(xs), np.asarray(ys)
    # ims_bytes = ims_bytes.reshape(sh)
    # mks_bytes = mks_bytes.reshape(sh)
    Hs = Hs.reshape(sh)  # should all be the same?
    Ws = Ws.reshape(sh)  # should all be the same?
    Ks = Ks.reshape(sh + (3, 3))  # should all be the same?
    xs = xs.reshape(sh)  # should all be the same?
    ys = ys.reshape(sh)  # should all be the same?

    return ims_bytes, mks_bytes, Ks, Hs, Ws, xs, ys
//  [python]
def load_resize_undist_ims_bytes(ims: np.ndarray,
                                 Ks: np.ndarray,
                                 Ds: np.ndarray,
                                 ratio: Union[float, List[int], List[float]] = 1.0,
                                 center_crop_size: List[int] = [-1, -1],
                                 desc="Loading image bytes from disk",
                                 **kwargs):
    sh = ims.shape  # V, N
    # Ks = np.broadcast_to(Ks[:, None], (sh + (3, 3)))
    # Ds = np.broadcast_to(Ds[:, None], (sh + (1, 5)))

    ims = ims.reshape((np.prod(sh)))
    # from easyvolcap.utils.dist_utils import get_rank
    # if not get_rank(): __import__('easyvolcap.utils.console_utils', fromlist=['debugger']).debugger()
    # else:
    #     while 1: pass
    Ks = Ks.reshape((np.prod(sh), 3, 3))
    Ds = Ds.reshape((np.prod(sh), 1, 5))

    ims = list(ims)
    Ks = list(Ks)
    Ds = list(Ds)  # only convert outer most dim to list

    if isinstance(ratio, list) and len(ratio) and isinstance(ratio[0], float):
        ratio = np.broadcast_to(np.asarray(ratio)[:, None], sh)  # V, N
        ratio = ratio.reshape((np.prod(sh)))
        ratio = list(ratio)
    elif isinstance(ratio, list):
        ratio = np.asarray(ratio)  # avoid expansion in parallel execution

    if isinstance(center_crop_size, list):
        center_crop_size = np.asarray(center_crop_size)  # avoid expansion

    # Should we batch these instead of loading?
    out = parallel_execution(ims, Ks, Ds, ratio, center_crop_size,
                             action=load_resize_undist_im_bytes,
                             desc=desc, print_progress=True,
                             **kwargs,
                             )

    ims_bytes, Ks, Hs, Ws = zip(*out)  # is this OK?
    ims_bytes, Ks, Hs, Ws = np.asarray(ims_bytes, dtype=object), np.asarray(Ks), np.asarray(Hs), np.asarray(Ws)
    # ims_bytes = ims_bytes.reshape(sh)  # numpy array of bytesio
    Hs = Hs.reshape(sh)  # should all be the same?
    Ws = Ws.reshape(sh)  # should all be the same?
    Ks = Ks.reshape(sh + (3, 3))  # should all be the same?

    return ims_bytes, Ks, Hs, Ws