Project: 012_Oscillatory_State-Space_Models | Function: apply_linoss_im (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Core implementation of LinOSS implicit discretization scheme solving the forced harmonic oscillator system using associative parallel scan

**KEY TERMS:** {
  "schur_comp": "S = (I + Δt²A/2)⁻¹ (Schur complement)",
  "M_IM_11": "M^{IM}_{11} matrix element",
  "M_IM_12": "M^{IM}_{12} matrix element",
  "M_IM_21": "M^{IM}_{21} matrix element",
  "M_IM_22": "M^{IM}_{22} matrix element",
  "associative_scan": "parallel scan O(log N) computation",
  "Bu_elements": "Bu_n input forcing terms"
}

**CONTEXT CODE:** from typing import List

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
from jax import nn
from jax.nn.initializers import normal
import math
from jax import random

def simple_uniform_init(rng, shape, std=1.):
    weights = random.uniform(rng, shape)*2.*std - std
    return weights

class GLU(eqx.Module):
    w1: eqx.nn.Linear
    w2: eqx.nn.Linear

    def __init__(self, input_dim, output_dim, key):
        w1_key, w2_key = jr.split(key, 2)
        self.w1 = eqx.nn.Linear(input_dim, output_dim, use_bias=True, key=w1_key)
        self.w2 = eqx.nn.Linear(input_dim, output_dim, use_bias=True, key=w2_key)

    def __call__(self, x):
        return self.w1(x) * jax.nn.sigmoid(self.w2(x))

# Parallel scan operations
@jax.vmap
def binary_operator(q_i, q_j):
    """ Binary operator for parallel scan of linear recurrence. Assumes a diagonal matrix A.
        Args:
            q_i: tuple containing A_i and Bu_i at position i       (P,), (P,)
            q_j: tuple containing A_j and Bu_j at position j       (P,), (P,)
        Returns:
            new element ( A_out, Bu_out )
    """
    raise NotImplementedError("Function implementation has been masked for testing")

def apply_linoss_im(A_diag, B, C_tilde, input_sequence, step):
    """Compute the LxH output of LinOSS-IM given an LxH input.
    Args:
        A_diag  (float32):   diagonal state matrix   (P,)
        B       (complex64): input matrix            (P, H)
        C       (complex64): output matrix           (H, P)
        input_sequence (float32): input sequence of features    (L, H)
        step    (float):     discretization time-step $\Delta_t$  (P,)
    Returns:
        outputs (float32): the SSM outputs (LinOSS_IMEX layer preactivations)      (L, H)
    """
    raise NotImplementedError("Function implementation has been masked for testing")

def apply_linoss_imex(A_diag, B, C, input_sequence, step):
    """Compute the LxH output of of LinOSS-IMEX given an LxH input.
    Args:
        A_diag  (float32):   diagonal state matrix   (P,)
        B       (complex64): input matrix            (P, H)
        C       (complex64): output matrix           (H, P)
        input_sequence (float32): input sequence of features    (L, H)
        step    (float):     discretization time-step $\Delta_t$  (P,)
    Returns:
        outputs (float32): the SSM outputs (LinOSS_IMEX layer preactivations)      (L, H)
    """
    Bu_elements = jax.vmap(lambda u: B @ u)(input_sequence)

    A_ = jnp.ones_like(A_diag)
    B_ = -1. * step * A_diag
    C_ = step * jnp.ones_like(A_diag)
    D_ = 1. - (step ** 2.) * A_diag

    M_IMEX = jnp.concatenate([A_, B_, C_, D_])

    M_IMEX_elements = M_IMEX * jnp.ones((input_sequence.shape[0],
                                          4 * A_diag.shape[0]))

    F1 = Bu_elements * step
    F2 = Bu_elements * (step ** 2.)
    F = jnp.hstack((F1, F2))

    _, xs = jax.lax.associative_scan(binary_operator, (M_IMEX_elements, F))
    ys = xs[:, A_diag.shape[0]:]

    return jax.vmap(lambda x: (C @ x).real)(ys)

class LinOSSLayer(eqx.Module):
    A_diag: jax.Array
    B: jax.Array
    C: jax.Array
    D: jax.Array
    steps: jax.Array
    discretization: str

    def __init__(
        self,
        ssm_size,
        H,
        discretization,
        *,
        key
    ):

        B_key, C_key, D_key, A_key, step_key, key = jr.split(key, 6)
        self.A_diag = random.uniform(A_key, shape=(ssm_size,))
        self.B = simple_uniform_init(B_key,shape=(ssm_size, H, 2),std=1./math.sqrt(H))
        self.C = simple_uniform_init(C_key,shape=(H, ssm_size, 2),std=1./math.sqrt(ssm_size))
        self.D = normal(stddev=1.0)(D_key, (H,))
        self.steps = random.uniform(step_key,shape=(ssm_size,))
        self.discretization = discretization

    def __call__(self, input_sequence):
        A_diag = nn.relu(self.A_diag)

        B_complex = self.B[..., 0] + 1j * self.B[..., 1]
        C_complex = self.C[..., 0] + 1j * self.C[..., 1]

        steps = nn.sigmoid(self.steps)
        if self.discretization == 'IMEX':
            ys = apply_linoss_imex(A_diag, B_complex, C_complex, input_sequence, steps)
        elif self.discretization == 'IM':
            ys = apply_linoss_im(A_diag, B_complex, C_complex, input_sequence, steps)
        else:
            print('Discretization type not implemented')

        Du = jax.vmap(lambda u: self.D * u)(input_sequence)
        return ys + Du


class LinOSSBlock(eqx.Module):

    norm: eqx.nn.BatchNorm
    ssm: LinOSSLayer
    glu: GLU
    drop: eqx.nn.Dropout

    def __init__(
        self,
        ssm_size,
        H,
        discretization,


Complete the function implementation replacing "raise NotImplementedError".