<commit_before>import collections
from itertools import repeat
from typing import Optional, Sequence, Union

import jax.numpy as jnp
import jax.random as jrandom
from jax.lax import conv_general_dilated

from ..custom_types import Array
from ..module import Module, static_field


def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable):
            if len(x) == n:
                return tuple(x)
            else:
                raise ValueError(f"Length of {x} (length = {len(x)}) is not equal to {n}")
        return tuple(repeat(x, n))

    return parse


class Conv(Module):
    num_spatial_dims: int = static_field()
    weight: Array
    bias: Optional[Array]
    in_channels: int = static_field()
    out_channels: int = static_field()
    kernel_size: Union[int, Sequence[int]] = static_field()
    stride: Union[int, Sequence[int]] = static_field()
    padding: Union[int, Sequence[int]] = static_field()
    dilation: Union[int, Sequence[int]] = static_field()
    use_bias: bool = static_field()

    def __init__(
        self,
        num_spatial_dims,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        dilation=1,
        use_bias=True,
        *,
        key,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.num_spatial_dims = num_spatial_dims
        parse = _ntuple(self.num_spatial_dims)
        wkey, bkey = jrandom.split(key, 2)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = parse(kernel_size)
        self.use_bias = use_bias
        lim = 1 / jnp.sqrt(self.in_channels * jnp.prod(jnp.array(self.kernel_size))).item()

        self.weight = jrandom.uniform(
            wkey,
            (
                self.out_channels,
                self.in_channels,
            )
            + self.kernel_size,
            minval=-lim,<commit_msg>remove .item()<commit_after>import collections
from itertools import repeat
from typing import Optional, Sequence, Union

import jax.numpy as jnp
import jax.random as jrandom
from jax.lax import conv_general_dilated

from ..custom_types import Array
from ..module import Module, static_field


def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable):
            if len(x) == n:
                return tuple(x)
            else:
                raise ValueError(f"Length of {x} (length = {len(x)}) is not equal to {n}")
        return tuple(repeat(x, n))

    return parse


class Conv(Module):
    num_spatial_dims: int = static_field()
    weight: Array
    bias: Optional[Array]
    in_channels: int = static_field()
    out_channels: int = static_field()
    kernel_size: Union[int, Sequence[int]] = static_field()
    stride: Union[int, Sequence[int]] = static_field()
    padding: Union[int, Sequence[int]] = static_field()
    dilation: Union[int, Sequence[int]] = static_field()
    use_bias: bool = static_field()

    def __init__(
        self,
        num_spatial_dims,
        in_channels,
        out_channels,
        kernel_size,
        stride=1,
        padding=0,
        dilation=1,
        use_bias=True,
        *,
        key,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.num_spatial_dims = num_spatial_dims
        parse = _ntuple(self.num_spatial_dims)
        wkey, bkey = jrandom.split(key, 2)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = parse(kernel_size)
        self.use_bias = use_bias
        lim = 1 / jnp.sqrt(self.in_channels * jnp.prod(jnp.array(self.kernel_size)))

        self.weight = jrandom.uniform(
            wkey,
            (
                self.out_channels,
                self.in_channels,
            )
            + self.kernel_size,
            minval=-lim,
{'fraction_non_alphanumeric': 0.10175619834710743, 'fraction_numerical': 0.0025826446280991736, 'mean_word_length': 2.238294314381271, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4395528', 'n_tokens_mistral': 1229, 'n_tokens_neox': 1195, 'n_words': 274}