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
import numpy as np
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_clone_0(in_ptr0, out_ptr0, 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 + x0, xmask)
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_clone_rsub_1(in_ptr0, in_ptr1, 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 + x0, xmask)
    tmp1 = tl.load(in_ptr1 + 0)
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK])
    tmp3 = tmp0 + tmp2
    tmp4 = tl.sigmoid(tmp3)
    tmp5 = 1.0
    tmp6 = tmp5 - tmp4
    tmp7 = tl.full([1], 0, tl.int32)
    tmp8 = tl.full([1], 1, tl.int32)
    tmp9 = tmp7 == tmp8
    tmp10 = tmp5 * tmp4
    tmp11 = tl.where(tmp9, tmp10, tmp5)
    tl.store(out_ptr0 + x0, tmp6, xmask)
    tl.store(out_ptr1 + x0, tmp11, xmask)


@triton.jit
def triton_poi_fused_clone_copy_mul_ones_2(in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 12
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 3
    x1 = xindex // 3
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr2 + x1, xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr3 + 0)
    tmp10 = tl.broadcast_to(tmp9, [XBLOCK])
    tmp0 = x0
    tmp1 = tl.full([1], 2, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp5 = tmp3 * tmp4
    tmp6 = tl.full([1], 1, tl.int32)
    tmp7 = tmp0 == tmp6
    tmp11 = tmp8 + tmp10
    tmp12 = tl.sigmoid(tmp11)
    tmp13 = 1.0
    tmp14 = tmp13 * tmp12
    tmp15 = tl.where(tmp7, tmp14, tmp13)
    tmp16 = tl.where(tmp2, tmp5, tmp15)
    tl.store(out_ptr0 + x2, tmp16, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4 = args
    args.clear()
    assert_size_stride(primals_1, (4, 1), (1, 4))
    assert_size_stride(primals_2, (4, 4), (4, 1))
    assert_size_stride(primals_3, (1,), (1,))
    assert_size_stride(primals_4, (4, 2), (2, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 1), (1, 4), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_clone_0[grid(4)](primals_1, buf0, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        del primals_1
        buf1 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_2, buf0, out=buf1)
        buf2 = reinterpret_tensor(buf0, (4,), (1,), 0)
        del buf0
        buf3 = empty_strided_cuda((4,), (1,), torch.float32)
        triton_poi_fused_clone_rsub_1[grid(4)](buf1, primals_3, buf2, buf3,
            4, XBLOCK=4, num_warps=1, num_stages=1)
        buf4 = empty_strided_cuda((4, 3), (3, 1), torch.float32)
        triton_poi_fused_clone_copy_mul_ones_2[grid(12)](buf3, buf2, buf1,
            primals_3, buf4, 12, XBLOCK=16, num_warps=1, num_stages=1)
        buf5 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        extern_kernels.mm(primals_4, reinterpret_tensor(buf4, (2, 4), (1, 3
            ), 1), out=buf5)
    return reinterpret_tensor(buf5, (4, 4), (1, 4), 0
        ), primals_3, buf1, buf2, buf3, reinterpret_tensor(primals_4, (2, 4
        ), (1, 2), 0), reinterpret_tensor(buf4, (4, 2), (3, 1), 1
        ), reinterpret_tensor(primals_2, (4, 4), (1, 4), 0)


class SoftTreeNew(torch.nn.Module):
    """Soft decision tree."""

    def __init__(self, in_features, out_features, depth, projection=
        'constant', dropout=0.0):
        super(SoftTreeNew, self).__init__()
        self.proj = projection
        self.depth = depth
        self.in_features = in_features
        self.out_features = out_features
        self.leaf_count = int(2 ** depth)
        self.gate_count = int(self.leaf_count - 1)
        self.gw = torch.nn.Parameter(torch.nn.init.kaiming_normal_(torch.
            empty(self.gate_count, in_features), nonlinearity='sigmoid').T)
        self.gb = torch.nn.Parameter(torch.zeros(self.gate_count))
        self.drop = torch.nn.Dropout(p=dropout)
        if self.proj == 'linear':
            self.pw = torch.nn.init.kaiming_normal_(torch.empty(
                out_features * self.leaf_count, in_features), nonlinearity=
                'linear')
            self.pw = torch.nn.Parameter(self.pw.view(out_features, self.
                leaf_count, in_features).permute(0, 2, 1))
            self.pb = torch.nn.Parameter(torch.zeros(out_features, self.
                leaf_count))
        elif self.proj == 'linear2':
            self.pw = torch.nn.init.kaiming_normal_(torch.empty(
                out_features * self.leaf_count, in_features), nonlinearity=
                'linear')
            self.pw = torch.nn.Parameter(self.pw.view(out_features, self.
                leaf_count, in_features).permute(1, 2, 0))
            self.pb = torch.nn.Parameter(torch.zeros(self.leaf_count, 1,
                out_features))
        elif self.proj == 'constant':
            self.z = torch.nn.Parameter(torch.randn(out_features, self.
                leaf_count))

    def extra_repr(self):
        return 'in_features=%d, out_features=%d, depth=%d, projection=%s' % (
            self.in_features, self.out_features, self.depth, self.proj)

    def node_densities(self, x):
        gw_ = self.drop(self.gw)
        gatings = torch.sigmoid(torch.add(torch.matmul(x, gw_), self.gb))
        node_densities = torch.ones(x.shape[0], 2 ** (self.depth + 1) - 1,
            device=x.device)
        it = 1
        for d in range(1, self.depth + 1):
            for i in range(2 ** d):
                parent_index = (it + 1) // 2 - 1
                child_way = (it + 1) % 2
                if child_way == 0:
                    parent_gating = gatings[:, parent_index]
                else:
                    parent_gating = 1 - gatings[:, parent_index]
                parent_density = node_densities[:, parent_index].clone()
                node_densities[:, it] = parent_density * parent_gating
                it += 1
        return node_densities

    def gatings(self, x):
        return torch.sigmoid(torch.add(torch.matmul(x, self.gw), self.gb))

    def total_path_value(self, x, index, level=None):
        gatings = self.gatings(x)
        gateways = np.binary_repr(index, width=self.depth)
        L = 0.0
        current = 0
        if level is None:
            level = self.depth
        for i in range(level):
            if int(gateways[i]) == 0:
                L += gatings[:, current].mean()
                current = 2 * current + 1
            else:
                L += (1 - gatings[:, current]).mean()
                current = 2 * current + 2
        return L

    def forward(self, input_0):
        primals_1 = self.gw
        primals_3 = self.gb
        primals_4 = self.z
        primals_2 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4])
        return output[0]


Measured latency_ms: 0.059647