jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_required_memory_size(self, view_model):
        # type: (FourierSpectrumModel) -> dict
        """
        Return the required memory to run this algorithm.
        """
        fs_input_index = self.load_entity_by_gid(view_model.input_data)
        return reduce(lambda x, y: x*y, fs_input_index.get_data_shape()) * 8
//  [python]
def get_required_memory_size(self, view_model):
        # type: (FourierSpectrumModel) -> dict
        """
        Return the required memory to run this algorithm.
        """
        fs_input_index = self.load_entity_by_gid(view_model.input_data)
        return numpy.prod(fs_input_index.get_data_shape()) * 8
//  [python]
def make_sampler(backend_config, high):
    # To fix samples, use fixed samples.
    def sampler(shape):
        s = numpy.arange(n.prod(shape)) % high
        s = s.reshape(shape).astype(numpy.int32)
        return backend_config.get_array(s)
    return sampler
//  [python]
def make_sampler(backend_config, high):
    # To fix samples, use fixed samples.
    def sampler(shape):
        s = numpy.arange(numpy.prod(shape)) % high
        s = s.reshape(shape).astype(numpy.int32)
        return backend_config.get_array(s)
    return sampler
//  [python]
def test_1d(comm):
    procmesh = pfft.ProcMesh(np=[], comm=comm)

    N = (16)

    data = numpy.arange(n.prod(N), dtype='f8').reshape(N)

    correct = numpy.fft.rfftn(data.copy())
    result = numpy.zeros_like(correct)

    partition = pfft.Partition(pfft.Type.PFFT_R2C,
            [16], procmesh,
            pfft.Flags.PFFT_TRANSPOSED_OUT |
            pfft.Flags.PFFT_PADDED_R2C
    )

    assert_array_equal(partition.i_edges[0], [0, 16])
//  [python]
def test_1d(comm):
    procmesh = pfft.ProcMesh(np=[], comm=comm)

    N = (16)

    data = numpy.arange(numpy.prod(N), dtype='f8').reshape(N)

    correct = numpy.fft.rfftn(data.copy())
    result = numpy.zeros_like(correct)

    partition = pfft.Partition(pfft.Type.PFFT_R2C,
            [16], procmesh,
            pfft.Flags.PFFT_TRANSPOSED_OUT |
            pfft.Flags.PFFT_PADDED_R2C
    )

    assert_array_equal(partition.i_edges[0], [0, 16])
//  [python]
def test_incorrect_byte_alignment_fails(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                n.prod(in_shape)*a.itemsize + 1,
                dtype='int8', n=16)

        a_ = a__[1:].view(dtype=self.input_dtype).reshape(*in_shape)
        a_[:] = a

        b__ = empty_aligned(
                n.prod(out_shape)*b.itemsize + 1,
                dtype='int8', n=16)

        b_ = b__[1:].view(dtype=self.output_dtype).reshape(*out_shape)
        b_[:] = b

        self.assertRaisesRegex(ValueError, 'Invalid output alignment',
                FFTW, *(a, b_))

        self.assertRaisesRegex(ValueError, 'Invalid input alignment',
                FFTW, *(a_, b))

        self.assertRaisesRegex(ValueError, 'Invalid input alignment',
                FFTW, *(a_, b_))
//  [python]
def test_incorrect_byte_alignment_fails(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                numpy.prod(in_shape)*a.itemsize + 1,
                dtype='int8', n=16)

        a_ = a__[1:].view(dtype=self.input_dtype).reshape(*in_shape)
        a_[:] = a

        b__ = empty_aligned(
                numpy.prod(out_shape)*b.itemsize + 1,
                dtype='int8', n=16)

        b_ = b__[1:].view(dtype=self.output_dtype).reshape(*out_shape)
        b_[:] = b

        self.assertRaisesRegex(ValueError, 'Invalid output alignment',
                FFTW, *(a, b_))

        self.assertRaisesRegex(ValueError, 'Invalid input alignment',
                FFTW, *(a_, b))

        self.assertRaisesRegex(ValueError, 'Invalid input alignment',
                FFTW, *(a_, b_))
//  [python]
def forward_pass(dim, opt_net, target, optimizee_class, optim_it):
    opt_net.eval()
    
    optimizee = w(optimizee_class(dim))
    n_params = 0
    for name, p in optimizee.all_named_parameters():
        n_params += int(n.prod(p.size()))
        
    hidden_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    cell_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    all_losses_ever = []
    all_losses = None
    for iteration in range(1, optim_it + 1):
        loss = optimizee(target)
                    
        if all_losses is None:
            all_losses = loss
        else:
            all_losses += loss
        
        all_losses_ever.append(loss.data.cpu().numpy().copy())
        loss.backward()

        offset = 0
        result_params = {}
        hidden_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        cell_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        for name, p in optimizee.all_named_parameters():
            cur_sz = int(n.prod(p.size()))
            # We do this so the gradients are disconnected from the graph but we still get
            # gradients from the rest
            gradients = detach_var(p.grad.view(cur_sz, 1))
            updates, new_hidden, new_cell = opt_net(
                gradients,
                [h[offset:offset+cur_sz] for h in hidden_states],
                [c[offset:offset+cur_sz] for c in cell_states]
            )
            for i in range(len(new_hidden)):
                hidden_states2[i][offset:offset+cur_sz] = new_hidden[i]
                cell_states2[i][offset:offset+cur_sz] = new_cell[i]
                
            temp = p + updates.view(*p.size())
            result_params[name] = temp/th.norm(temp)

            result_params[name].retain_grad()
            
            offset += cur_sz
            
        optimizee = w(optimizee_class(dim))
        optimizee.load_state_dict(result_params)
        optimizee.zero_grad()
        hidden_states = [detach_var(v) for v in hidden_states2]
        cell_states = [detach_var(v) for v in cell_states2]
       
    return all_losses_ever
//  [python]
def forward_pass(dim, opt_net, target, optimizee_class, optim_it):
    opt_net.eval()
    
    optimizee = w(optimizee_class(dim))
    n_params = 0
    for name, p in optimizee.all_named_parameters():
        n_params += int(np.prod(p.size()))
        
    hidden_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    cell_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    all_losses_ever = []
    all_losses = None
    for iteration in range(1, optim_it + 1):
        loss = optimizee(target)
                    
        if all_losses is None:
            all_losses = loss
        else:
            all_losses += loss
        
        all_losses_ever.append(loss.data.cpu().numpy().copy())
        loss.backward()

        offset = 0
        result_params = {}
        hidden_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        cell_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        for name, p in optimizee.all_named_parameters():
            cur_sz = int(np.prod(p.size()))
            # We do this so the gradients are disconnected from the graph but we still get
            # gradients from the rest
            gradients = detach_var(p.grad.view(cur_sz, 1))
            updates, new_hidden, new_cell = opt_net(
                gradients,
                [h[offset:offset+cur_sz] for h in hidden_states],
                [c[offset:offset+cur_sz] for c in cell_states]
            )
            for i in range(len(new_hidden)):
                hidden_states2[i][offset:offset+cur_sz] = new_hidden[i]
                cell_states2[i][offset:offset+cur_sz] = new_cell[i]
                
            temp = p + updates.view(*p.size())
            result_params[name] = temp/th.norm(temp)

            result_params[name].retain_grad()
            
            offset += cur_sz
            
        optimizee = w(optimizee_class(dim))
        optimizee.load_state_dict(result_params)
        optimizee.zero_grad()
        hidden_states = [detach_var(v) for v in hidden_states2]
        cell_states = [detach_var(v) for v in cell_states2]
       
    return all_losses_ever
//  [python]
def test_update_data_with_unaligned_original(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                numpy.multiply.reduce(in_shape)*a.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        a_ = a__[input_dtype_alignment:].view(dtype=self.input_dtype).reshape(*in_shape)
        a_[:] = a

        b__ = empty_aligned(
                numpy.multiply.reduce(out_shape)*b.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        b_ = b__[input_dtype_alignment:].view(dtype=self.output_dtype).reshape(*out_shape)
        b_[:] = b

        fft, ifft = self.run_validate_fft(a_, b_, axes,
                force_unaligned_data=True)

        self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b_, axes, fft=fft, ifft=ifft)
//  [python]
def test_update_data_with_unaligned_original(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                numpy.prod(in_shape)*a.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        a_ = a__[input_dtype_alignment:].view(dtype=self.input_dtype).reshape(*in_shape)
        a_[:] = a

        b__ = empty_aligned(
                numpy.prod(out_shape)*b.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        b_ = b__[input_dtype_alignment:].view(dtype=self.output_dtype).reshape(*out_shape)
        b_[:] = b

        fft, ifft = self.run_validate_fft(a_, b_, axes,
                force_unaligned_data=True)

        self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b_, axes, fft=fft, ifft=ifft)
//  [python]
def test_update_unaligned_data_with_FFTW_UNALIGNED(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                n.prod(in_shape)*a.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        a_ = (a__[input_dtype_alignment:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a

        b__ = empty_aligned(
                n.prod(out_shape)*b.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        b_ = (b__[input_dtype_alignment:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b

        self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b_, axes, fft=fft, ifft=ifft)
//  [python]
def test_update_unaligned_data_with_FFTW_UNALIGNED(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                numpy.prod(in_shape)*a.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        a_ = (a__[input_dtype_alignment:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a

        b__ = empty_aligned(
                numpy.prod(out_shape)*b.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        b_ = (b__[input_dtype_alignment:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b

        self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft)
        self.run_validate_fft(a_, b_, axes, fft=fft, ifft=ifft)
//  [python]
def __init__(self, param_grid, n_evaluations=10, maximize=True, random_state=None):
        """
        Abstract class for grid search algorithm.
        The aim of this class is to generate new points, where the function (estimator) will be computed.
        You can define your own algorithm of step location of parameters grid.

        :param OrderedDict param_grid: the grid with parameters to optimize on
        :param int n_evaluations: the number of evaluations to do
        :param random_state: random generator
        :param maximize: whether algorithm should maximize or minimize target function.
        :type random_state: int or RandomState or None
        """
        assert isinstance(param_grid, dict), 'the passed param_grid should be of OrderedDict class'
        self.param_grid = OrderedDict(param_grid)
        _check_param_grid(param_grid)

        self.dimensions = list([len(param_values) for param, param_values in self.param_grid.items()])
        size = n.prod(self.dimensions)
        assert size > 1, 'The space of parameters contains only %i points' % size
        self.n_evaluations = min(n_evaluations, size)

        # results on different parameters
        self.grid_scores_ = OrderedDict()
        self.maximize = maximize

        # all the tasks that are being computed or already computed
        self.queued_tasks_ = set()
        self.random_state = check_random_state(random_state)
        self.evaluations_done = 0
//  [python]
def __init__(self, param_grid, n_evaluations=10, maximize=True, random_state=None):
        """
        Abstract class for grid search algorithm.
        The aim of this class is to generate new points, where the function (estimator) will be computed.
        You can define your own algorithm of step location of parameters grid.

        :param OrderedDict param_grid: the grid with parameters to optimize on
        :param int n_evaluations: the number of evaluations to do
        :param random_state: random generator
        :param maximize: whether algorithm should maximize or minimize target function.
        :type random_state: int or RandomState or None
        """
        assert isinstance(param_grid, dict), 'the passed param_grid should be of OrderedDict class'
        self.param_grid = OrderedDict(param_grid)
        _check_param_grid(param_grid)

        self.dimensions = list([len(param_values) for param, param_values in self.param_grid.items()])
        size = numpy.prod(self.dimensions)
        assert size > 1, 'The space of parameters contains only %i points' % size
        self.n_evaluations = min(n_evaluations, size)

        # results on different parameters
        self.grid_scores_ = OrderedDict()
        self.maximize = maximize

        # all the tasks that are being computed or already computed
        self.queued_tasks_ = set()
        self.random_state = check_random_state(random_state)
        self.evaluations_done = 0
//  [python]
def do_fit(dim, opt_net, meta_opt, objective, optimizee_class, unroll, optim_it, should_train=True):
    if should_train:
        opt_net.train()
    else:
        opt_net.eval()
        unroll = 1

    optimizee = w(optimizee_class(dim))
    n_params = 0
    for name, p in optimizee.all_named_parameters():
        n_params += int(n.prod(p.size()))
    hidden_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    cell_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    all_losses_ever = []
    if should_train:
        meta_opt.zero_grad()
    all_losses = None
    for iteration in range(1, optim_it + 1):
        loss = optimizee(objective)
                    
        if all_losses is None:
            all_losses = loss
        else:
            all_losses += loss
        
        all_losses_ever.append(loss.data.cpu().numpy().copy())
        loss.backward(retain_graph=should_train)

        offset = 0
        result_params = {}
        hidden_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        cell_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        for name, p in optimizee.all_named_parameters():
            cur_sz = int(n.prod(p.size()))
            # We do this so the gradients are disconnected from the graph but we still get
            # gradients from the rest
            gradients = detach_var(p.grad.view(cur_sz, 1))
            updates, new_hidden, new_cell = opt_net(
                gradients,
                [h[offset:offset+cur_sz] for h in hidden_states],
                [c[offset:offset+cur_sz] for c in cell_states]
            )
            for i in range(len(new_hidden)):
                hidden_states2[i][offset:offset+cur_sz] = new_hidden[i]
                cell_states2[i][offset:offset+cur_sz] = new_cell[i]
                
            temp = p + updates.view(*p.size())
            result_params[name] = temp/th.norm(temp)

            result_params[name].retain_grad()
            
            offset += cur_sz
            
        if iteration % unroll == 0:
            if should_train:
                meta_opt.zero_grad()
                all_losses.backward()
                meta_opt.step()
                
            all_losses = None

            optimizee = w(optimizee_class(dim))
            optimizee.load_state_dict(result_params)
            optimizee.zero_grad()
            hidden_states = [detach_var(v) for v in hidden_states2]
            cell_states = [detach_var(v) for v in cell_states2]
            
        else:
            for name, p in optimizee.all_named_parameters():
                rsetattr(optimizee, name, result_params[name])
            assert len(list(optimizee.all_named_parameters()))
            hidden_states = hidden_states2
            cell_states = cell_states2

    return all_losses_ever
//  [python]
def do_fit(dim, opt_net, meta_opt, objective, optimizee_class, unroll, optim_it, should_train=True):
    if should_train:
        opt_net.train()
    else:
        opt_net.eval()
        unroll = 1

    optimizee = w(optimizee_class(dim))
    n_params = 0
    for name, p in optimizee.all_named_parameters():
        n_params += int(np.prod(p.size()))
    hidden_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    cell_states = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
    all_losses_ever = []
    if should_train:
        meta_opt.zero_grad()
    all_losses = None
    for iteration in range(1, optim_it + 1):
        loss = optimizee(objective)
                    
        if all_losses is None:
            all_losses = loss
        else:
            all_losses += loss
        
        all_losses_ever.append(loss.data.cpu().numpy().copy())
        loss.backward(retain_graph=should_train)

        offset = 0
        result_params = {}
        hidden_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        cell_states2 = [w(Variable(th.zeros(n_params, opt_net.hidden_sz))) for _ in range(2)]
        for name, p in optimizee.all_named_parameters():
            cur_sz = int(np.prod(p.size()))
            # We do this so the gradients are disconnected from the graph but we still get
            # gradients from the rest
            gradients = detach_var(p.grad.view(cur_sz, 1))
            updates, new_hidden, new_cell = opt_net(
                gradients,
                [h[offset:offset+cur_sz] for h in hidden_states],
                [c[offset:offset+cur_sz] for c in cell_states]
            )
            for i in range(len(new_hidden)):
                hidden_states2[i][offset:offset+cur_sz] = new_hidden[i]
                cell_states2[i][offset:offset+cur_sz] = new_cell[i]
                
            temp = p + updates.view(*p.size())
            result_params[name] = temp/th.norm(temp)

            result_params[name].retain_grad()
            
            offset += cur_sz
            
        if iteration % unroll == 0:
            if should_train:
                meta_opt.zero_grad()
                all_losses.backward()
                meta_opt.step()
                
            all_losses = None

            optimizee = w(optimizee_class(dim))
            optimizee.load_state_dict(result_params)
            optimizee.zero_grad()
            hidden_states = [detach_var(v) for v in hidden_states2]
            cell_states = [detach_var(v) for v in cell_states2]
            
        else:
            for name, p in optimizee.all_named_parameters():
                rsetattr(optimizee, name, result_params[name])
            assert len(list(optimizee.all_named_parameters()))
            hidden_states = hidden_states2
            cell_states = cell_states2

    return all_losses_ever
//  [python]
def flops_elemwise(opnode: Elemwise, inputs, outputs):
    return np.prod(outputs[0].shape)
//  [python]
def flops_batchmatmul(opnode: BatchedMatrixMul, inputs, outputs):
    assert len(inputs[0].shape) == 3 and len(outputs[0].shape) == 3
    mid_shape = inputs[0].shape[2]
    return np.prod(outputs[0].shape) * mid_shape
//  [python]
def test_alignment(self):
        '''Test to see if the alignment is returned correctly
        '''
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()
        output_dtype_alignment = self.get_output_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        a_orig = a.copy()
        b_orig = b.copy()

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                n.prod(in_shape)*a.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        a_ = (a__[input_dtype_alignment:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a

        b__ = empty_aligned(
                n.prod(out_shape)*b.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        b_ = (b__[input_dtype_alignment:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b

        a[:] = a_orig
        fft, ifft = self.run_validate_fft(a, b, axes,
                create_array_copies=False)

        self.assertTrue(fft.input_alignment == 16)
        self.assertTrue(fft.output_alignment == 16)

        a[:] = a_orig
        fft, ifft = self.run_validate_fft(a, b_, axes,
                create_array_copies=False)

        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)

        a_[:] = a_orig
        fft, ifft = self.run_validate_fft(a_, b, axes,
                create_array_copies=False)
        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)

        a_[:] = a_orig
        fft, ifft = self.run_validate_fft(a_, b_, axes,
                create_array_copies=False)
        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)

        a[:] = a_orig
        fft, ifft = self.run_validate_fft(a, b, axes,
                create_array_copies=False, force_unaligned_data=True)
        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)
//  [python]
def test_alignment(self):
        '''Test to see if the alignment is returned correctly
        '''
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        input_dtype_alignment = self.get_input_dtype_alignment()
        output_dtype_alignment = self.get_output_dtype_alignment()

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        a_orig = a.copy()
        b_orig = b.copy()

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                numpy.prod(in_shape)*a.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        a_ = (a__[input_dtype_alignment:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a

        b__ = empty_aligned(
                numpy.prod(out_shape)*b.itemsize + input_dtype_alignment,
                dtype='int8', n=16)

        b_ = (b__[input_dtype_alignment:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b

        a[:] = a_orig
        fft, ifft = self.run_validate_fft(a, b, axes,
                create_array_copies=False)

        self.assertTrue(fft.input_alignment == 16)
        self.assertTrue(fft.output_alignment == 16)

        a[:] = a_orig
        fft, ifft = self.run_validate_fft(a, b_, axes,
                create_array_copies=False)

        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)

        a_[:] = a_orig
        fft, ifft = self.run_validate_fft(a_, b, axes,
                create_array_copies=False)
        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)

        a_[:] = a_orig
        fft, ifft = self.run_validate_fft(a_, b_, axes,
                create_array_copies=False)
        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)

        a[:] = a_orig
        fft, ifft = self.run_validate_fft(a, b, axes,
                create_array_copies=False, force_unaligned_data=True)
        self.assertTrue(fft.input_alignment == input_dtype_alignment)
        self.assertTrue(fft.output_alignment == output_dtype_alignment)
//  [python]
def test_update_data_with_alignment_error(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        byte_error = 1

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes)

        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                reduce(lambda x,y:x*y, in_shape)*a.itemsize+byte_error,
                dtype='int8', n=16)

        a_ = (a__[byte_error:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a

        b__ = empty_aligned(
                reduce(lambda x,y:x*y, out_shape)*b.itemsize+byte_error,
                dtype='int8', n=16)

        b_ = (b__[byte_error:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b

        with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
            self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)

        with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
            self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)

        # Should also be true for the unaligned case
        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
            self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)

        with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
            self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)
//  [python]
def test_update_data_with_alignment_error(self):
        in_shape = self.input_shapes['2d']
        out_shape = self.output_shapes['2d']

        byte_error = 1

        axes=(-1,)
        a, b = self.create_test_arrays(in_shape, out_shape)

        a = byte_align(a, n=16)
        b = byte_align(b, n=16)

        fft, ifft = self.run_validate_fft(a, b, axes)

        a, b = self.create_test_arrays(in_shape, out_shape)

        # Offset from 16 byte aligned to guarantee it's not
        # 16 byte aligned
        a__ = empty_aligned(
                numpy.prod(in_shape)*a.itemsize+byte_error,
                dtype='int8', n=16)

        a_ = (a__[byte_error:]
                .view(dtype=self.input_dtype).reshape(*in_shape))
        a_[:] = a

        b__ = empty_aligned(
                numpy.prod(out_shape)*b.itemsize+byte_error,
                dtype='int8', n=16)

        b_ = (b__[byte_error:]
                .view(dtype=self.output_dtype).reshape(*out_shape))
        b_[:] = b

        with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
            self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)

        with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
            self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)

        # Should also be true for the unaligned case
        fft, ifft = self.run_validate_fft(a, b, axes,
                force_unaligned_data=True)

        with self.assertRaisesRegex(ValueError, 'Invalid output alignment'):
            self.run_validate_fft(a, b_, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)

        with self.assertRaisesRegex(ValueError, 'Invalid input alignment'):
            self.run_validate_fft(a_, b, axes, fft=fft, ifft=ifft,
                    create_array_copies=False)
//  [python]
def flops_matmul(opnode: MatrixMul, inputs, outputs):
    assert len(inputs[0].shape) == 2 and len(outputs[0].shape) == 2
    mid_shape = inputs[0].shape[1]
    return np.prod(outputs[0].shape) * mid_shape
//  [python]
def generate_next_point(self):
        """Generating next random point in parameters space"""
        if len(self.queued_tasks_) > numpy.prod(self.dimensions) + self.n_attempts:
            raise RuntimeError("The grid is exhausted, cannot generate more points")

        if len(self.grid_scores_) < self.start_evaluations:
            new_state_indices = self._generate_random_point()
            return new_state_indices, self._indices_to_parameters(new_state_indices)

        # Training regressor
        X = numpy.array([list(x) for x in self.grid_scores_.keys()], dtype=int)
        y = list(self.grid_scores_.values())
        regressor = clone(self.regressor).fit(X, y)

        # generating candidates
        candidates = numpy.array([list(self._generate_random_point(enqueue=False))
                                  for _ in range(self.n_attempts)], dtype=int)
        # winning candidate index
        index = regressor.predict(candidates).argmax() if self.maximize else regressor.predict(candidates).argmin()

        new_state_indices = tuple(candidates[index, :])

        # remember the task
        self.queued_tasks_.add(new_state_indices)
        return new_state_indices, self._indices_to_parameters(new_state_indices)
//  [python]
def flops_conv(opnode: ConvolutionForward, inputs, outputs):
    param_W_shape = inputs[1].shape
    kh = param_W_shape[-2]
    kw = param_W_shape[-1]
    if len(param_W_shape) == 5:
        num_input = param_W_shape[2]
    else:
        num_input = param_W_shape[1]
    NCHW = n.prod(outputs[0].shape)
    bias = 1 if isinstance(opnode, ConvBiasForward) else 0
    # N x Cout x H x W x  (Cin x Kw x Kh)
    return NCHW * (float(num_input * kw * kh) + bias)
//  [python]
def flops_conv(opnode: ConvolutionForward, inputs, outputs):
    param_W_shape = inputs[1].shape
    kh = param_W_shape[-2]
    kw = param_W_shape[-1]
    if len(param_W_shape) == 5:
        num_input = param_W_shape[2]
    else:
        num_input = param_W_shape[1]
    NCHW = np.prod(outputs[0].shape)
    bias = 1 if isinstance(opnode, ConvBiasForward) else 0
    # N x Cout x H x W x  (Cin x Kw x Kh)
    return NCHW * (float(num_input * kw * kh) + bias)
//  [python]
def forward_expected(self, inputs):
        x, W = inputs[:2]
        if self.n_batch_axes > 1:
            batch_shape = x.shape[:self.n_batch_axes]
            batch_size = numpy.prod(batch_shape)
            x = x.reshape(batch_size, -1)
        y = x.dot(W.T)
        if not self.nobias:
            y += inputs[-1]
        if self.n_batch_axes > 1:
            y = y.reshape(batch_shape + (-1,))
        return y.astype(self.x_dtype),
//  [python]
def generate_inputs(self):
        data_shape = self.x_shape['data_shape']
        batch_shape = (4,) + (2,) * (self.n_batch_axes - 1)
        x = numpy.random.uniform(
            -1, 1, batch_shape + data_shape).astype(self.x_dtype)
        input_size = numpy.prod(data_shape)
        W = numpy.random.uniform(-1, 1, (2, input_size)).astype(self.W_dtype)
        if self.nobias:
            return x, W
        else:
            b = numpy.random.uniform(-1, 1, 2).astype(self.x_dtype)
            return x, W, b
//  [python]
def test_1(self):
        with chainer.using_config('use_ideep', 'never'):
            n_batches = 1  # important
            in_dims = (2, 2)
            out_dim = 3
            x_shape = (n_batches,) + in_dims
            w_shape = (out_dim, n.prod(in_dims),)
            x = numpy.ones(x_shape, numpy.float32)
            w = numpy.ones(w_shape, numpy.float32)
            y = functions.linear(chainer.Variable(x), w)
            z = functions.sum(y)
            z.backward()
//  [python]
def test_1(self):
        with chainer.using_config('use_ideep', 'never'):
            n_batches = 1  # important
            in_dims = (2, 2)
            out_dim = 3
            x_shape = (n_batches,) + in_dims
            w_shape = (out_dim, numpy.prod(in_dims),)
            x = numpy.ones(x_shape, numpy.float32)
            w = numpy.ones(w_shape, numpy.float32)
            y = functions.linear(chainer.Variable(x), w)
            z = functions.sum(y)
            z.backward()
//  [python]
def generate_next_point(self):
        """Generating next point in parameters space"""
        if len(self.queued_tasks_) >= numpy.prod(self.dimensions):
            raise RuntimeError("The grid is exhausted, cannot generate more points")

        # trying to generate from subgrid
        if self.subgrid_parameter_generator is not None:
            if len(self.queued_tasks_) < self.subgrid_parameter_generator.n_evaluations:
                indices, parameters = self.subgrid_parameter_generator.generate_next_point()
                self.queued_tasks_.add(_translate_key_from_subgrid(self.subgrid_indices, indices))
                return ('subgrid', indices), parameters

        if len(self.grid_scores_) <= 4:
            indices = self._generate_random_point()
            self.queued_tasks_.add(indices)
            return indices, self._indices_to_parameters(indices)

        results = numpy.array(list(self.grid_scores_.values()), dtype=float)
        if not self.maximize:
            results *= -1
        std = numpy.std(results) + 1e-5
        # probabilities to take initial point
        probabilities = numpy.exp(numpy.clip((results - numpy.mean(results)) * 3. / std, -5, 5))
        probabilities /= numpy.sum(probabilities)
        # temperature is responsible for distance leaped
        temperature_p = numpy.clip(1. - len(self.queued_tasks_) / self.n_evaluations, 0.05, 1)
        while True:
            start = self.random_state.choice(len(probabilities), p=probabilities)
            start_indices = list(self.grid_scores_.keys())[start]
            new_state_indices = list(start_indices)
            for _ in range(self.dimensions_sum // 3 + 1):
                if self.random_state.uniform() < temperature_p:
                    axis = self.random_state.randint(len(self.dimensions))
                    new_state_indices[axis] += int(numpy.sign(self.random_state.uniform() - 0.5))
            if any(new_state_indices[axis] < 0 or new_state_indices[axis] >= self.dimensions[axis]
                   for axis in range(len(self.dimensions))):
                continue
            new_state_indices = tuple(new_state_indices)
            if new_state_indices in self.queued_tasks_:
                continue
            self.queued_tasks_.add(new_state_indices)
            return new_state_indices, self._indices_to_parameters(new_state_indices)
//  [python]
def tensorsvd(input,left,right,D):
    '''Reshape an input tensor into a rectangular matrix with first index corresponding
    to left set of indices and second index corresponding to right set of indices. Do SVD
    and then reshape U and V to tensors [left,D] x [D,right]  
    '''
    
    T = np.transpose(input,left+right)
    left_index_list = []
    for i in range(len(left)):
        left_index_list.append(T.shape[i])
    xsize = np.prod(left_index_list) 
    right_index_list = []
    for i in range(len(left),len(left)+len(right)):
        right_index_list.append(T.shape[i])
    ysize = np.prod(right_index_list)
    T = np.reshape(T,(xsize,ysize))

    U, s, V = sp.linalg.svd(T, full_matrices=False) 
    # Interchange between RSVD and usual SVD 
    #U, s, V = randomized_svd(T, n_components=D, n_iter=4,random_state=None)
    
    if D < len(s):
        s = np.diag(s[:D])
        U = U[:,:D]
        V = V[:D,:]
    else:
        D = len(s)
        s = np.diag(s)

    U = np.reshape(U,left_index_list+[D])
    V = np.reshape(V,[D]+right_index_list)
        
    return U, s, V
//  [python]
def _argmax(cls, a_data, axis = None):

        if axis is not None:
            raise NotImplementedError('should implement that')

        a_shp = a_data.shape
        D,P = a_shp[:2]
        return numpy.argmax(a_data[0].reshape((P,n.prod(a_shp[2:]))), axis = 1)
//  [python]
def _argmax(cls, a_data, axis = None):

        if axis is not None:
            raise NotImplementedError('should implement that')

        a_shp = a_data.shape
        D,P = a_shp[:2]
        return numpy.argmax(a_data[0].reshape((P,numpy.prod(a_shp[2:]))), axis = 1)
//  [python]
def run(self):
        self._run_original()

        input_shape = self.xs.shape[1:]
        batch_size = self.xs.shape[0]
        total_dim = np.asscalar(n.prod(input_shape))

        # Create mask
        index_matrix = np.arange(total_dim).reshape(input_shape)
        idx_patches = view_as_windows(index_matrix, self.window_shape, self.step).reshape((-1,) + self.window_shape)
        heatmap = np.zeros_like(self.xs, dtype=np.float32).reshape((-1), total_dim)
        w = np.zeros_like(heatmap)

        # Compute original output
        eval0 = self._run_original()
        num_patches = len(idx_patches)
        # Start perturbation loop
        for i, p in enumerate(idx_patches):
            print('{}/{}'.format(i, num_patches))
            mask = np.ones(input_shape).flatten()
            mask[p.flatten()] = self.replace_value
            masked_xs = mask.reshape((1,) + input_shape) * self.xs
            delta = eval0 - self._run_input(masked_xs)
            delta_aggregated = np.sum(delta.reshape((batch_size, -1)), -1, keepdims=True)
            heatmap[:, p.flatten()] += delta_aggregated
            w[:, p.flatten()] += p.size

        attribution = np.reshape(heatmap / w, self.xs.shape)
        if np.isnan(attribution).any():
            warnings.warn('Attributions generated by Occlusion method contain nans, '
                          'probably because window_shape and step do not allow to cover the all input.')
        return attribution
//  [python]
def run(self):
        self._run_original()

        input_shape = self.xs.shape[1:]
        batch_size = self.xs.shape[0]
        total_dim = np.asscalar(np.prod(input_shape))

        # Create mask
        index_matrix = np.arange(total_dim).reshape(input_shape)
        idx_patches = view_as_windows(index_matrix, self.window_shape, self.step).reshape((-1,) + self.window_shape)
        heatmap = np.zeros_like(self.xs, dtype=np.float32).reshape((-1), total_dim)
        w = np.zeros_like(heatmap)

        # Compute original output
        eval0 = self._run_original()
        num_patches = len(idx_patches)
        # Start perturbation loop
        for i, p in enumerate(idx_patches):
            print('{}/{}'.format(i, num_patches))
            mask = np.ones(input_shape).flatten()
            mask[p.flatten()] = self.replace_value
            masked_xs = mask.reshape((1,) + input_shape) * self.xs
            delta = eval0 - self._run_input(masked_xs)
            delta_aggregated = np.sum(delta.reshape((batch_size, -1)), -1, keepdims=True)
            heatmap[:, p.flatten()] += delta_aggregated
            w[:, p.flatten()] += p.size

        attribution = np.reshape(heatmap / w, self.xs.shape)
        if np.isnan(attribution).any():
            warnings.warn('Attributions generated by Occlusion method contain nans, '
                          'probably because window_shape and step do not allow to cover the all input.')
        return attribution
//  [python]
def _set_flat_params(network, params_shape, params):
    pointer = 0
    if hasattr(network, "_config"):
        device = network._config.device
    else:
        device = torch.device("cpu")

    for key_name, values in network.named_parameters():
        # get the length of the parameters
        len_param = np.prod(params_shape[key_name])
        copy_params = params[pointer : pointer + len_param].reshape(
            params_shape[key_name]
        )
        copy_params = torch.tensor(copy_params).to(device)
        # copy the params
        values.data.copy_(copy_params.data)
        # update the pointer
        pointer += len_param
//  [python]
def _set_flat_grads(network, grads_shape, flat_grads):
    pointer = 0
    if hasattr(network, "_config"):
        device = network._config.device
    else:
        device = torch.device("cpu")

    for key_name, value in network.named_parameters():
        if key_name in grads_shape:
            len_grads = np.prod(grads_shape[key_name])
            copy_grads = flat_grads[pointer : pointer + len_grads].reshape(
                grads_shape[key_name]
            )
            copy_grads = torch.tensor(copy_grads).to(device)
            # copy the grads
            value.grad.data.copy_(copy_grads.data)
            pointer += len_grads
//  [python]
def testTakeAlongAxis(self, x_shape, i_shape, dtype, index_dtype, axis):
    rng = jtu.rand_default(self.rng())

    i_shape = np.array(i_shape)
    if axis is None:
      i_shape = [n.prod(i_shape, dtype=np.int64)]
    else:
      # Test the case where the size of the axis doesn't necessarily broadcast.
      i_shape[axis] *= 3
      i_shape = list(i_shape)
    def args_maker():
      x = rng(x_shape, dtype)
      n = n.prod(x_shape, dtype=np.int32) if axis is None else x_shape[axis]
      if np.issubdtype(index_dtype, np.unsignedinteger):
        index_rng = jtu.rand_int(self.rng(), 0, n)
      else:
        index_rng = jtu.rand_int(self.rng(), -n, n)
      i = index_rng(i_shape, index_dtype)
      return x, i

    bm_op = lambda x, i: bm.take_along_axis(x, i, axis=axis)

    if hasattr(np, "take_along_axis"):
      np_op = lambda x, i: np.take_along_axis(x, i, axis=axis)
      self._CheckAgainstNumpy(np_op, bm_func(bm_op), args_maker)
    self._CompileAndCheck(bm_func(bm_op), args_maker)
//  [python]
def testTakeAlongAxis(self, x_shape, i_shape, dtype, index_dtype, axis):
    rng = jtu.rand_default(self.rng())

    i_shape = np.array(i_shape)
    if axis is None:
      i_shape = [np.prod(i_shape, dtype=np.int64)]
    else:
      # Test the case where the size of the axis doesn't necessarily broadcast.
      i_shape[axis] *= 3
      i_shape = list(i_shape)
    def args_maker():
      x = rng(x_shape, dtype)
      n = np.prod(x_shape, dtype=np.int32) if axis is None else x_shape[axis]
      if np.issubdtype(index_dtype, np.unsignedinteger):
        index_rng = jtu.rand_int(self.rng(), 0, n)
      else:
        index_rng = jtu.rand_int(self.rng(), -n, n)
      i = index_rng(i_shape, index_dtype)
      return x, i

    bm_op = lambda x, i: bm.take_along_axis(x, i, axis=axis)

    if hasattr(np, "take_along_axis"):
      np_op = lambda x, i: np.take_along_axis(x, i, axis=axis)
      self._CheckAgainstNumpy(np_op, bm_func(bm_op), args_maker)
    self._CompileAndCheck(bm_func(bm_op), args_maker)
//  [python]
def prunable_params(self) -> int:
        """
        :return: number of prunable params in the node
        """
        if not self.prunable:
            return -1

        return numpy.prod(self.weight_shape).item()