# AOT ID: ['0_forward']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, grid_combo_kernels, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()


# kernel path: runs/run_shard_7/inductor_cache/mz/cmzggcufrudhcxbs54wcdqp6m6pkyknmtvr6y5vrlwnrrksqtcye.py
# Topologically Sorted Source Nodes: [j], Original ATen: [aten._softmax]
# Source node to ATen node mapping:
#   j => amax, exp, sub, sum_1
# Graph fragment:
#   %amax : [num_users=1] = call_function[target=torch.ops.aten.amax.default](args = (%unsqueeze, [-1], True), kwargs = {})
#   %sub : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%unsqueeze, %amax), kwargs = {})
#   %exp : [num_users=2] = call_function[target=torch.ops.aten.exp.default](args = (%sub,), kwargs = {})
#   %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%exp, [-1], True), kwargs = {})
triton_poi_fused__softmax_0 = async_compile.triton('triton_poi_fused__softmax_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[4], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__softmax_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused__softmax_0(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK : tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (5*x0), xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr0 + (1 + (5*x0)), xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (2 + (5*x0)), xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (3 + (5*x0)), xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr0 + (4 + (5*x0)), xmask, eviction_policy='evict_last')
    tmp2 = triton_helpers.maximum(tmp0, tmp1)
    tmp4 = triton_helpers.maximum(tmp2, tmp3)
    tmp6 = triton_helpers.maximum(tmp4, tmp5)
    tmp8 = triton_helpers.maximum(tmp6, tmp7)
    tmp9 = tmp0 - tmp8
    tmp10 = tl_math.exp(tmp9)
    tmp11 = tmp1 - tmp8
    tmp12 = tl_math.exp(tmp11)
    tmp13 = tmp10 + tmp12
    tmp14 = tmp3 - tmp8
    tmp15 = tl_math.exp(tmp14)
    tmp16 = tmp13 + tmp15
    tmp17 = tmp5 - tmp8
    tmp18 = tl_math.exp(tmp17)
    tmp19 = tmp16 + tmp18
    tmp20 = tmp7 - tmp8
    tmp21 = tl_math.exp(tmp20)
    tmp22 = tmp19 + tmp21
    tl.store(out_ptr0 + (x0), tmp8, xmask)
    tl.store(out_ptr1 + (x0), tmp22, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/hp/chpgaxjmfea3iltrobasajt6dfuvw23vqc72imv4luy3xwlhmj5l.py
# Topologically Sorted Source Nodes: [j], Original ATen: [aten._softmax]
# Source node to ATen node mapping:
#   j => amax, div, exp, sub, sum_1
# Graph fragment:
#   %amax : [num_users=1] = call_function[target=torch.ops.aten.amax.default](args = (%unsqueeze, [-1], True), kwargs = {})
#   %sub : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%unsqueeze, %amax), kwargs = {})
#   %exp : [num_users=2] = call_function[target=torch.ops.aten.exp.default](args = (%sub,), kwargs = {})
#   %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%exp, [-1], True), kwargs = {})
#   %div : [num_users=2] = call_function[target=torch.ops.aten.div.Tensor](args = (%exp, %sum_1), kwargs = {})
triton_poi_fused__softmax_1 = async_compile.triton('triton_poi_fused__softmax_1', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[32], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__softmax_1', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 3, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 20
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = (xindex // 5)
    tmp0 = tl.load(in_ptr0 + (x2), xmask)
    tmp1 = tl.load(in_ptr1 + (x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr2 + (x1), xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp5 = tmp3 / tmp4
    tl.store(out_ptr0 + (x2), tmp5, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/77/c77fzlbbirm4s52xcl6wixpwu6fb5paaithoa547xxck3vwqoi3l.py
# Topologically Sorted Source Nodes: [r_1], Original ATen: [aten.view]
# Source node to ATen node mapping:
#   r_1 => full_default
# Graph fragment:
#   %full_default : [num_users=2] = call_function[target=torch.ops.aten.full.default](args = ([1, 1], 1.0), kwargs = {dtype: torch.float32, layout: torch.strided, device: cuda:0, pin_memory: False})
triton_poi_fused_view_2 = async_compile.triton('triton_poi_fused_view_2', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[1], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {1: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0,), equal_to_1=(1,))]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_view_2', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 0, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_view_2(out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 1
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = tl.full([XBLOCK], True, tl.int1)
    tmp0 = 1.0
    tl.store(out_ptr0 + (tl.full([XBLOCK], 0, tl.int32)), tmp0, None)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/mz/cmz72f3f3w3wlapqxei3eeokxsml27mism6a47wuqnhqq4xvto5u.py
# Topologically Sorted Source Nodes: [matmul], Original ATen: [aten.view]
# Source node to ATen node mapping:
#   matmul => view_2
# Graph fragment:
#   %view_2 : [num_users=2] = call_function[target=torch.ops.aten.reshape.default](args = (%unsqueeze_1, [4, 4]), kwargs = {})
triton_poi_fused_view_3 = async_compile.triton('triton_poi_fused_view_3', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_view_3', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 1, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_view_3(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask)
    tmp1 = tl_math.log(tmp0)
    tl.store(out_ptr0 + (x0), tmp1, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/j2/cj24z5wiqbewlsbpgaxnwgd6fvcvx2jl37z2yhhzlgeznoniu7ww.py
# Topologically Sorted Source Nodes: [add, exp], Original ATen: [aten.add, aten.exp]
# Source node to ATen node mapping:
#   add => add
#   exp => exp_1
# Graph fragment:
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%slice_2, %primals_5), kwargs = {})
#   %exp_1 : [num_users=1] = call_function[target=torch.ops.aten.exp.default](args = (%add,), kwargs = {})
triton_poi_fused_add_exp_4 = async_compile.triton('triton_poi_fused_add_exp_4', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[16], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_exp_4', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 3, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_add_exp_4(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = (xindex // 4)
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (x0), xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr1 + (x0 + (5*x1)), xmask)
    tmp4 = tl.load(in_ptr2 + (x0), xmask, eviction_policy='evict_last')
    tmp1 = tl.sigmoid(tmp0)
    tmp3 = tmp1 * tmp2
    tmp5 = tmp3 + tmp4
    tmp6 = tl_math.exp(tmp5)
    tl.store(out_ptr0 + (x2), tmp6, xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (4, 5), (5, 1))
    assert_size_stride(primals_3, (5, 1), (1, 1))
    assert_size_stride(primals_4, (5, ), (1, ))
    assert_size_stride(primals_5, (4, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 4, 1), (4, 1, 4), torch.float32)
        buf1 = empty_strided_cuda((1, 4, 1), (4, 1, 4), torch.float32)
        # Topologically Sorted Source Nodes: [j], Original ATen: [aten._softmax]
        stream0 = get_raw_stream(0)
        triton_poi_fused__softmax_0.run(primals_2, buf0, buf1, 4, grid=grid(4), stream=stream0)
        buf2 = empty_strided_cuda((1, 4, 5), (20, 5, 1), torch.float32)
        # Topologically Sorted Source Nodes: [j], Original ATen: [aten._softmax]
        triton_poi_fused__softmax_1.run(primals_2, buf0, buf1, buf2, 20, grid=grid(20), stream=stream0)
        del buf0
        del buf1
        del primals_2
        buf3 = empty_strided_cuda((1, 1), (1, 1), torch.float32)
        # Topologically Sorted Source Nodes: [r_1], Original ATen: [aten.view]
        triton_poi_fused_view_2.run(buf3, 1, grid=grid(1), stream=stream0)
        buf4 = empty_strided_cuda((1, 5), (5, 1), torch.float32)
        # Topologically Sorted Source Nodes: [r_1], Original ATen: [aten.addmm]
        extern_kernels.addmm(primals_4, buf3, reinterpret_tensor(primals_3, (1, 5), (1, 1), 0), alpha=1, beta=1, out=buf4)
        del primals_3
        del primals_4
        buf5 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [matmul], Original ATen: [aten.view]
        triton_poi_fused_view_3.run(primals_1, buf5, 16, grid=grid(16), stream=stream0)
        del primals_1
        buf6 = empty_strided_cuda((4, 5), (5, 1), torch.float32)
        # Topologically Sorted Source Nodes: [matmul], Original ATen: [aten.mm]
        extern_kernels.mm(buf5, reinterpret_tensor(buf2, (4, 5), (5, 1), 0), out=buf6)
        buf7 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [add, exp], Original ATen: [aten.add, aten.exp]
        triton_poi_fused_add_exp_4.run(buf4, buf6, primals_5, buf7, 16, grid=grid(16), stream=stream0)
        del primals_5
    return (buf7, buf2, buf3, buf4, buf6, buf7, reinterpret_tensor(buf5, (4, 4), (1, 4), 0), )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    primals_1 = rand_strided((4, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_2 = rand_strided((4, 5), (5, 1), device='cuda:0', dtype=torch.float32)
    primals_3 = rand_strided((5, 1), (1, 1), device='cuda:0', dtype=torch.float32)
    primals_4 = rand_strided((5, ), (1, ), device='cuda:0', dtype=torch.float32)
    primals_5 = rand_strided((4, ), (1, ), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([primals_1, primals_2, primals_3, primals_4, primals_5])
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    compiled_module_main('None', benchmark_compiled_module)

import collections
import torch
import torch.utils.data
from torch import nn


def get_redistribution(kind: 'str', num_states: 'int', num_features: 'int'=
    None, num_out: 'int'=None, normaliser: 'nn.Module'=None, **kwargs):
    if kind == 'linear':
        return LinearRedistribution(num_states, num_features, num_out,
            normaliser)
    elif kind == 'outer':
        return OuterRedistribution(num_states, num_features, num_out,
            normaliser, **kwargs)
    elif kind == 'gate':
        return GateRedistribution(num_states, num_features, num_out, normaliser
            )
    else:
        raise ValueError('unknown kind of redistribution: {}'.format(kind))


class SummaryWriterNamespaceNoLoggingScope:

    def __init__(self, writer):
        self._writer = writer

    def __enter__(self):
        self._writer._logging_enabled = False

    def __exit__(self, type, value, traceback):
        self._writer._logging_enabled = True
        return False


class DummySummaryWriter:

    def __init__(self, **kwargs):
        self._logging_enabled = False
        pass

    def add_scalar(self, name, value, verbose_only=True):
        pass

    def add_summary(self, name, tensor, verbose_only=True):
        pass

    def add_histogram(self, name, tensor, verbose_only=True):
        pass

    def add_tensor(self, name, tensor, verbose_only=True):
        pass

    def print(self, name, tensor, verbose_only=True):
        pass

    def namespace(self, name):
        return self

    def every(self, epoch_interval):
        return self

    def verbose(self, verbose):
        return self

    def no_logging(self):
        return SummaryWriterNamespaceNoLoggingScope(self)


class NoRandomScope:

    def __init__(self, module):
        self._module = module

    def __enter__(self):
        self._module._disable_random()

    def __exit__(self, type, value, traceback):
        self._module._enable_random()
        return False


class ExtendedTorchModule(torch.nn.Module):

    def __init__(self, default_name, *args, writer=None, name=None, **kwargs):
        super().__init__()
        if writer is None:
            writer = DummySummaryWriter()
        self.writer = writer.namespace(default_name if name is None else name)
        self.allow_random = True

    def set_parameter(self, name, value):
        parameter = getattr(self, name, None)
        if isinstance(parameter, torch.nn.Parameter):
            parameter.fill_(value)
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module.set_parameter(name, value)

    def regualizer(self, merge_in=None):
        regualizers = collections.defaultdict(int)
        if merge_in is not None:
            for key, value in merge_in.items():
                self.writer.add_scalar(f'regualizer/{key}', value)
                regualizers[key] += value
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                for key, value in module.regualizer().items():
                    regualizers[key] += value
        return regualizers

    def optimize(self, loss):
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module.optimize(loss)

    def log_gradients(self):
        for name, parameter in self.named_parameters(recurse=False):
            if parameter.requires_grad:
                gradient, *_ = parameter.grad.data
                self.writer.add_summary(f'{name}/grad', gradient)
                self.writer.add_histogram(f'{name}/grad', gradient)
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module.log_gradients()

    def no_internal_logging(self):
        return self.writer.no_logging()

    def _disable_random(self):
        self.allow_random = False
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module._disable_random()

    def _enable_random(self):
        self.allow_random = True
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module._enable_random()

    def no_random(self):
        return NoRandomScope(self)


class NormalisedSigmoid(nn.Module):
    """ Normalised logistic sigmoid function. """

    def __init__(self, p: 'float'=1, dim: 'int'=-1):
        super().__init__()
        self.p = p
        self.dim = dim

    def forward(self, s: 'torch.Tensor') ->torch.Tensor:
        a = torch.sigmoid(s)
        return torch.nn.functional.normalize(a, p=self.p, dim=self.dim)


class Redistribution(nn.Module):
    """ Base class for modules that generate redistribution vectors/matrices. """

    def __init__(self, num_states: 'int', num_features: 'int'=None, num_out:
        'int'=None, normaliser: 'nn.Module'=None):
        """
        Parameters
        ----------
        num_states : int
            The number of states this redistribution is to be applied on.
        num_features : int, optional
            The number of features to use for configuring the redistribution.
            If the redistribution is not input-dependent, this argument will be ignored.
        num_out : int, optional
            The number of outputs to redistribute the states to.
            If nothing is specified, the redistribution matrix is assumed to be square.
        normaliser : Module, optional
            Function to use for normalising the redistribution matrix.
        """
        super().__init__()
        self.num_features = num_features
        self.num_states = num_states
        self.num_out = num_out or num_states
        self.normaliser = normaliser or NormalisedSigmoid(dim=-1)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        raise NotImplementedError('subclass must implement this method')

    def forward(self, x: 'torch.Tensor') ->torch.Tensor:
        r = self._compute(x)
        return self.normaliser(r)


class Gate(Redistribution):
    """
    Classic gate as used in e.g. LSTMs.

    Notes
    -----
    The vector that is computed by this module gives rise to a diagonal redistribution matrix,
    i.e. a redistribution matrix that does not really redistribute (not normalised).
    """

    def __init__(self, num_states, num_features, num_out=None, sigmoid=None):
        super().__init__(num_states, num_features, 1, sigmoid or nn.Sigmoid())
        self.fc = nn.Linear(num_features, num_states)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.orthogonal_(self.fc.weight)
        nn.init.zeros_(self.fc.bias)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        return self.fc(x)


class GateRedistribution(Redistribution):
    """
    Gate-like redistribution that only depends on input.

    This module directly computes all entries for the redistribution matrix
    from a linear combination of the input values and is normalised by the activation function.
    """

    def __init__(self, num_states, num_features, num_out=None, normaliser=None
        ):
        super().__init__(num_states, num_features, num_out, normaliser)
        self.fc = nn.Linear(num_features, self.num_states * self.num_out)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.orthogonal_(self.fc.weight)
        nn.init.zeros_(self.fc.bias)
        if self.num_states == self.num_out:
            with torch.no_grad():
                self.fc.bias[0::self.num_out + 1] = 3

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        logits = self.fc(x)
        return logits.view(-1, self.num_states, self.num_out)


class LinearRedistribution(Redistribution):
    """
    Redistribution by normalising a learned matrix.

    This module has an unnormalised version of the redistribution matrix as parameters
    and is normalised by applying a non-linearity (the normaliser).
    The redistribution does not depend on any of the input values,
    but is updated using the gradients to fit the data.
    """

    def __init__(self, num_states, num_features=0, num_out=None, normaliser
        =None):
        super(LinearRedistribution, self).__init__(num_states, 0, num_out,
            normaliser)
        self.r = nn.Parameter(torch.empty(self.num_states, self.num_out),
            requires_grad=True)
        self.reset_parameters()

    @torch.no_grad()
    def reset_parameters(self):
        if self.num_states == self.num_out:
            nn.init.eye_(self.r)
            if type(self.normaliser) is NormalisedSigmoid:
                torch.mul(self.r, 2, out=self.r)
                torch.sub(self.r, 1, out=self.r)
        else:
            nn.init.orthogonal_(self.r)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        return self.r.unsqueeze(0)


class OuterRedistribution(Redistribution):
    """
    Redistribution by (weighted) outer product of two input-dependent vectors.

    This module computes the entries for the redistribution matrix as
    the outer product of two vectors that are linear combinations of the input values.
    There is an option to include a weight matrix parameter
    to weight each entry in the resulting matrix, which is then normalised using a non-linearity.
    The weight matrix parameter is updated through the gradients to fit the data.
    """

    def __init__(self, num_states, num_features, num_out=None, normaliser=
        None, weighted: 'bool'=False):
        """
        Parameters
        ----------
        weighted : bool, optional
            Whether or not to use a weighted outer product.
        """
        super(OuterRedistribution, self).__init__(num_states, num_features,
            num_out, normaliser)
        self.weighted = weighted
        self.r = nn.Parameter(torch.empty(self.num_states, self.num_out),
            requires_grad=weighted)
        self.fc1 = nn.Linear(num_features, self.num_states)
        self.fc2 = nn.Linear(num_features, self.num_out)
        self.phi = lambda x: x
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.ones_(self.r)
        nn.init.orthogonal_(self.fc1.weight)
        nn.init.zeros_(self.fc1.bias)
        nn.init.orthogonal_(self.fc2.weight)
        nn.init.zeros_(self.fc2.bias)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        a1 = self.phi(self.fc1(x))
        a2 = self.phi(self.fc2(x))
        outer = a1.unsqueeze(-1) * a2.unsqueeze(-2)
        if self.weighted:
            outer *= self.r
        return outer


class MCFullyConnected(ExtendedTorchModule):

    def __init__(self, in_features: 'int', out_features: 'int', **kwargs):
        super().__init__('MCFC', **kwargs)
        self.mass_input_size = in_features
        self.aux_input_size = 1
        self.hidden_size = out_features
        self.normaliser = nn.Softmax(dim=-1)
        self.out_gate = Gate(self.hidden_size, self.aux_input_size)
        self.junction = get_redistribution('linear', num_states=self.
            mass_input_size, num_features=self.aux_input_size, num_out=self
            .hidden_size, normaliser=self.normaliser)

    @torch.no_grad()
    def reset_parameters(self):
        self.out_gate.reset_parameters()
        self.junction.reset_parameters()

    def log_gradients(self):
        for name, parameter in self.named_parameters():
            gradient, *_ = parameter.grad.data
            self.writer.add_summary(f'{name}/grad', gradient)
            self.writer.add_histogram(f'{name}/grad', gradient)

    def regualizer(self, merge_in=None):
        r1 = -torch.mean(self.junction.r ** 2)
        r2 = -torch.mean(self.out_gate.fc.weight ** 2)
        r3 = -torch.mean(self.out_gate.fc.bias ** 2)
        return super().regualizer({'W': r1 + r2 + r3})

    def forward(self, x):
        x_m, x_a = x, x.new_ones(1)
        j = self.junction(x_a)
        o = self.out_gate(x_a)
        m_in = torch.matmul(x_m.unsqueeze(-2), j).squeeze(-2)
        return o * m_in


class MulMCFC(ExtendedTorchModule):

    def __init__(self, in_features, out_features, **kwargs):
        super().__init__('MulMCFC', **kwargs)
        self.mcfc = MCFullyConnected(in_features, out_features + 1, **kwargs)
        self.bias = nn.Parameter(torch.zeros(out_features))

    def reset_parameters(self):
        self.mcfc.reset_parameters()
        nn.init.zeros_(self.bias)

    def forward(self, x):
        log_sum = self.mcfc(torch.log(x))[:, :-1]
        return torch.exp(log_sum + self.bias)


def get_inputs():
    return [torch.rand([4, 4])]


def get_init_inputs():
    return [[], {'in_features': 4, 'out_features': 4}]

import torch
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import collections
import torch.utils.data
from torch import nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused__softmax_0(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr0 + (1 + 5 * x0), xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (2 + 5 * x0), xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (3 + 5 * x0), xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr0 + (4 + 5 * x0), xmask, eviction_policy='evict_last')
    tmp2 = triton_helpers.maximum(tmp0, tmp1)
    tmp4 = triton_helpers.maximum(tmp2, tmp3)
    tmp6 = triton_helpers.maximum(tmp4, tmp5)
    tmp8 = triton_helpers.maximum(tmp6, tmp7)
    tmp9 = tmp0 - tmp8
    tmp10 = tl_math.exp(tmp9)
    tmp11 = tmp1 - tmp8
    tmp12 = tl_math.exp(tmp11)
    tmp13 = tmp10 + tmp12
    tmp14 = tmp3 - tmp8
    tmp15 = tl_math.exp(tmp14)
    tmp16 = tmp13 + tmp15
    tmp17 = tmp5 - tmp8
    tmp18 = tl_math.exp(tmp17)
    tmp19 = tmp16 + tmp18
    tmp20 = tmp7 - tmp8
    tmp21 = tl_math.exp(tmp20)
    tmp22 = tmp19 + tmp21
    tl.store(out_ptr0 + x0, tmp8, xmask)
    tl.store(out_ptr1 + x0, tmp22, xmask)


@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 20
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 5
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr2 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp5 = tmp3 / tmp4
    tl.store(out_ptr0 + x2, tmp5, xmask)


@triton.jit
def triton_poi_fused_view_2(out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    tmp0 = 1.0
    tl.store(out_ptr0 + tl.full([XBLOCK], 0, tl.int32), tmp0, None)


@triton.jit
def triton_poi_fused_view_3(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl_math.log(tmp0)
    tl.store(out_ptr0 + x0, tmp1, xmask)


@triton.jit
def triton_poi_fused_add_exp_4(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr1 + (x0 + 5 * x1), xmask)
    tmp4 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.sigmoid(tmp0)
    tmp3 = tmp1 * tmp2
    tmp5 = tmp3 + tmp4
    tmp6 = tl_math.exp(tmp5)
    tl.store(out_ptr0 + x2, tmp6, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (4, 5), (5, 1))
    assert_size_stride(primals_3, (5, 1), (1, 1))
    assert_size_stride(primals_4, (5,), (1,))
    assert_size_stride(primals_5, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 4, 1), (4, 1, 4), torch.float32)
        buf1 = empty_strided_cuda((1, 4, 1), (4, 1, 4), torch.float32)
        get_raw_stream(0)
        triton_poi_fused__softmax_0[grid(4)](primals_2, buf0, buf1, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf2 = empty_strided_cuda((1, 4, 5), (20, 5, 1), torch.float32)
        triton_poi_fused__softmax_1[grid(20)](primals_2, buf0, buf1, buf2, 
            20, XBLOCK=32, num_warps=1, num_stages=1)
        del buf0
        del buf1
        del primals_2
        buf3 = empty_strided_cuda((1, 1), (1, 1), torch.float32)
        triton_poi_fused_view_2[grid(1)](buf3, 1, XBLOCK=1, num_warps=1,
            num_stages=1)
        buf4 = empty_strided_cuda((1, 5), (5, 1), torch.float32)
        extern_kernels.addmm(primals_4, buf3, reinterpret_tensor(primals_3,
            (1, 5), (1, 1), 0), alpha=1, beta=1, out=buf4)
        del primals_3
        del primals_4
        buf5 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_view_3[grid(16)](primals_1, buf5, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del primals_1
        buf6 = empty_strided_cuda((4, 5), (5, 1), torch.float32)
        extern_kernels.mm(buf5, reinterpret_tensor(buf2, (4, 5), (5, 1), 0),
            out=buf6)
        buf7 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_add_exp_4[grid(16)](buf4, buf6, primals_5, buf7, 
            16, XBLOCK=16, num_warps=1, num_stages=1)
        del primals_5
    return buf7, buf2, buf3, buf4, buf6, buf7, reinterpret_tensor(buf5, (4,
        4), (1, 4), 0)


def get_redistribution(kind: 'str', num_states: 'int', num_features: 'int'=
    None, num_out: 'int'=None, normaliser: 'nn.Module'=None, **kwargs):
    if kind == 'linear':
        return LinearRedistribution(num_states, num_features, num_out,
            normaliser)
    elif kind == 'outer':
        return OuterRedistribution(num_states, num_features, num_out,
            normaliser, **kwargs)
    elif kind == 'gate':
        return GateRedistribution(num_states, num_features, num_out, normaliser
            )
    else:
        raise ValueError('unknown kind of redistribution: {}'.format(kind))


class SummaryWriterNamespaceNoLoggingScope:

    def __init__(self, writer):
        self._writer = writer

    def __enter__(self):
        self._writer._logging_enabled = False

    def __exit__(self, type, value, traceback):
        self._writer._logging_enabled = True
        return False


class DummySummaryWriter:

    def __init__(self, **kwargs):
        self._logging_enabled = False
        pass

    def add_scalar(self, name, value, verbose_only=True):
        pass

    def add_summary(self, name, tensor, verbose_only=True):
        pass

    def add_histogram(self, name, tensor, verbose_only=True):
        pass

    def add_tensor(self, name, tensor, verbose_only=True):
        pass

    def print(self, name, tensor, verbose_only=True):
        pass

    def namespace(self, name):
        return self

    def every(self, epoch_interval):
        return self

    def verbose(self, verbose):
        return self

    def no_logging(self):
        return SummaryWriterNamespaceNoLoggingScope(self)


class NoRandomScope:

    def __init__(self, module):
        self._module = module

    def __enter__(self):
        self._module._disable_random()

    def __exit__(self, type, value, traceback):
        self._module._enable_random()
        return False


class ExtendedTorchModule(torch.nn.Module):

    def __init__(self, default_name, *args, writer=None, name=None, **kwargs):
        super().__init__()
        if writer is None:
            writer = DummySummaryWriter()
        self.writer = writer.namespace(default_name if name is None else name)
        self.allow_random = True

    def set_parameter(self, name, value):
        parameter = getattr(self, name, None)
        if isinstance(parameter, torch.nn.Parameter):
            parameter.fill_(value)
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module.set_parameter(name, value)

    def regualizer(self, merge_in=None):
        regualizers = collections.defaultdict(int)
        if merge_in is not None:
            for key, value in merge_in.items():
                self.writer.add_scalar(f'regualizer/{key}', value)
                regualizers[key] += value
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                for key, value in module.regualizer().items():
                    regualizers[key] += value
        return regualizers

    def optimize(self, loss):
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module.optimize(loss)

    def log_gradients(self):
        for name, parameter in self.named_parameters(recurse=False):
            if parameter.requires_grad:
                gradient, *_ = parameter.grad.data
                self.writer.add_summary(f'{name}/grad', gradient)
                self.writer.add_histogram(f'{name}/grad', gradient)
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module.log_gradients()

    def no_internal_logging(self):
        return self.writer.no_logging()

    def _disable_random(self):
        self.allow_random = False
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module._disable_random()

    def _enable_random(self):
        self.allow_random = True
        for module in self.children():
            if isinstance(module, ExtendedTorchModule):
                module._enable_random()

    def no_random(self):
        return NoRandomScope(self)


class NormalisedSigmoid(nn.Module):
    """ Normalised logistic sigmoid function. """

    def __init__(self, p: 'float'=1, dim: 'int'=-1):
        super().__init__()
        self.p = p
        self.dim = dim

    def forward(self, s: 'torch.Tensor') ->torch.Tensor:
        a = torch.sigmoid(s)
        return torch.nn.functional.normalize(a, p=self.p, dim=self.dim)


class Redistribution(nn.Module):
    """ Base class for modules that generate redistribution vectors/matrices. """

    def __init__(self, num_states: 'int', num_features: 'int'=None, num_out:
        'int'=None, normaliser: 'nn.Module'=None):
        """
        Parameters
        ----------
        num_states : int
            The number of states this redistribution is to be applied on.
        num_features : int, optional
            The number of features to use for configuring the redistribution.
            If the redistribution is not input-dependent, this argument will be ignored.
        num_out : int, optional
            The number of outputs to redistribute the states to.
            If nothing is specified, the redistribution matrix is assumed to be square.
        normaliser : Module, optional
            Function to use for normalising the redistribution matrix.
        """
        super().__init__()
        self.num_features = num_features
        self.num_states = num_states
        self.num_out = num_out or num_states
        self.normaliser = normaliser or NormalisedSigmoid(dim=-1)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        raise NotImplementedError('subclass must implement this method')

    def forward(self, x: 'torch.Tensor') ->torch.Tensor:
        r = self._compute(x)
        return self.normaliser(r)


class Gate(Redistribution):
    """
    Classic gate as used in e.g. LSTMs.

    Notes
    -----
    The vector that is computed by this module gives rise to a diagonal redistribution matrix,
    i.e. a redistribution matrix that does not really redistribute (not normalised).
    """

    def __init__(self, num_states, num_features, num_out=None, sigmoid=None):
        super().__init__(num_states, num_features, 1, sigmoid or nn.Sigmoid())
        self.fc = nn.Linear(num_features, num_states)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.orthogonal_(self.fc.weight)
        nn.init.zeros_(self.fc.bias)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        return self.fc(x)


class GateRedistribution(Redistribution):
    """
    Gate-like redistribution that only depends on input.

    This module directly computes all entries for the redistribution matrix
    from a linear combination of the input values and is normalised by the activation function.
    """

    def __init__(self, num_states, num_features, num_out=None, normaliser=None
        ):
        super().__init__(num_states, num_features, num_out, normaliser)
        self.fc = nn.Linear(num_features, self.num_states * self.num_out)
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.orthogonal_(self.fc.weight)
        nn.init.zeros_(self.fc.bias)
        if self.num_states == self.num_out:
            with torch.no_grad():
                self.fc.bias[0::self.num_out + 1] = 3

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        logits = self.fc(x)
        return logits.view(-1, self.num_states, self.num_out)


class LinearRedistribution(Redistribution):
    """
    Redistribution by normalising a learned matrix.

    This module has an unnormalised version of the redistribution matrix as parameters
    and is normalised by applying a non-linearity (the normaliser).
    The redistribution does not depend on any of the input values,
    but is updated using the gradients to fit the data.
    """

    def __init__(self, num_states, num_features=0, num_out=None, normaliser
        =None):
        super(LinearRedistribution, self).__init__(num_states, 0, num_out,
            normaliser)
        self.r = nn.Parameter(torch.empty(self.num_states, self.num_out),
            requires_grad=True)
        self.reset_parameters()

    @torch.no_grad()
    def reset_parameters(self):
        if self.num_states == self.num_out:
            nn.init.eye_(self.r)
            if type(self.normaliser) is NormalisedSigmoid:
                torch.mul(self.r, 2, out=self.r)
                torch.sub(self.r, 1, out=self.r)
        else:
            nn.init.orthogonal_(self.r)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        return self.r.unsqueeze(0)


class OuterRedistribution(Redistribution):
    """
    Redistribution by (weighted) outer product of two input-dependent vectors.

    This module computes the entries for the redistribution matrix as
    the outer product of two vectors that are linear combinations of the input values.
    There is an option to include a weight matrix parameter
    to weight each entry in the resulting matrix, which is then normalised using a non-linearity.
    The weight matrix parameter is updated through the gradients to fit the data.
    """

    def __init__(self, num_states, num_features, num_out=None, normaliser=
        None, weighted: 'bool'=False):
        """
        Parameters
        ----------
        weighted : bool, optional
            Whether or not to use a weighted outer product.
        """
        super(OuterRedistribution, self).__init__(num_states, num_features,
            num_out, normaliser)
        self.weighted = weighted
        self.r = nn.Parameter(torch.empty(self.num_states, self.num_out),
            requires_grad=weighted)
        self.fc1 = nn.Linear(num_features, self.num_states)
        self.fc2 = nn.Linear(num_features, self.num_out)
        self.phi = lambda x: x
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.ones_(self.r)
        nn.init.orthogonal_(self.fc1.weight)
        nn.init.zeros_(self.fc1.bias)
        nn.init.orthogonal_(self.fc2.weight)
        nn.init.zeros_(self.fc2.bias)

    def _compute(self, x: 'torch.Tensor') ->torch.Tensor:
        a1 = self.phi(self.fc1(x))
        a2 = self.phi(self.fc2(x))
        outer = a1.unsqueeze(-1) * a2.unsqueeze(-2)
        if self.weighted:
            outer *= self.r
        return outer


class MCFullyConnected(ExtendedTorchModule):

    def __init__(self, in_features: 'int', out_features: 'int', **kwargs):
        super().__init__('MCFC', **kwargs)
        self.mass_input_size = in_features
        self.aux_input_size = 1
        self.hidden_size = out_features
        self.normaliser = nn.Softmax(dim=-1)
        self.out_gate = Gate(self.hidden_size, self.aux_input_size)
        self.junction = get_redistribution('linear', num_states=self.
            mass_input_size, num_features=self.aux_input_size, num_out=self
            .hidden_size, normaliser=self.normaliser)

    @torch.no_grad()
    def reset_parameters(self):
        self.out_gate.reset_parameters()
        self.junction.reset_parameters()

    def log_gradients(self):
        for name, parameter in self.named_parameters():
            gradient, *_ = parameter.grad.data
            self.writer.add_summary(f'{name}/grad', gradient)
            self.writer.add_histogram(f'{name}/grad', gradient)

    def regualizer(self, merge_in=None):
        r1 = -torch.mean(self.junction.r ** 2)
        r2 = -torch.mean(self.out_gate.fc.weight ** 2)
        r3 = -torch.mean(self.out_gate.fc.bias ** 2)
        return super().regualizer({'W': r1 + r2 + r3})

    def forward(self, x):
        x_m, x_a = x, x.new_ones(1)
        j = self.junction(x_a)
        o = self.out_gate(x_a)
        m_in = torch.matmul(x_m.unsqueeze(-2), j).squeeze(-2)
        return o * m_in


class MulMCFCNew(ExtendedTorchModule):

    def __init__(self, in_features, out_features, **kwargs):
        super().__init__('MulMCFC', **kwargs)
        self.mcfc = MCFullyConnected(in_features, out_features + 1, **kwargs)
        self.bias = nn.Parameter(torch.zeros(out_features))

    def reset_parameters(self):
        self.mcfc.reset_parameters()
        nn.init.zeros_(self.bias)

    def forward(self, input_0):
        primals_5 = self.bias
        primals_3 = self.mcfc.out_gate.fc.weight
        primals_4 = self.mcfc.out_gate.fc.bias
        primals_2 = self.mcfc.junction.r
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4, primals_5])
        return output[0]
