You are an expert in writing code in many different languages. Your goal is to perform code completion for the following code, keeping in mind the rest of the code and the file meta data.
    Metadata details:
    filename: select_assign.rs
    path of file: ./repos/burn/crates/burn-jit/src/kernel/index
    
    the code of the file until where you have to start completion: use crate::{
    codegen::{
        dialect::gpu::{gpu, Branch, Elem, Item, Scope, Variable, Visibility},
        execute_dynamic, Compilation, CompilationInfo, CompilationSettings, Compiler, EagerHandle,
        InputInfo, WorkgroupLaunch,
    },
    element::JitElement,
    kernel::{elemwise_workgroup, DynamicKernelSource, SourceTemplate, WORKGROUP_DEFAULT},
    tensor::JitTensor,
    Runtime,
};
use std::marker::PhantomData;

pub struct SelectAssignComputeShader {
    tensor: Variable,
    indices: Variable,
    value: Variable,
    dim: usize,
}

#[derive(new)]
pub struct SelectAssignEagerKernel<R: Runtime, E: JitElement> {
    dim: usize,
    _runtime: PhantomData<R>,
    _elem: PhantomData<E>,
}

impl SelectAssignComputeShader {
    pub fn expand(self, scope: &mut Scope) {
        let tensor = self.tensor;
        let value = self.value;
        let indices = self.indices;
        let id = Variable::Id;

        let offset_tensor = scope.zero(Elem::UInt);
        let offset_value = scope.zero(Elem::UInt);

        let stride_tensor_dim = scope.create_local(Elem::UInt);
        let stride_value_dim = scope.create_local(Elem::UInt);
        let shape_value_dim = scope.create_local(Elem::UInt);

        let num_elems = scope.create_local(Elem::UInt);
        gpu!(scope, num_elems = cast(1u32));

        gpu!(
            scope,
            range(0u32, Variable::Rank).for_each(|i, scope| {
                let shape_value = scope.create_local(Elem::UInt);
                let stride_tensor = scope.create_local(Elem::UInt);
                let stride_value = scope.create_local(Elem::UInt);

                gpu!(scope, stride_tensor = stride(tensor, i));
                gpu!(scope, stride_value = stride(value, i));

    
    pub fn expand(self, scope: &mut Scope) {
        let tensor = self.tensor;
        let value = self.value;
        let indices = self.indices;
        let id = Variable::Id;

        let offset_tensor = scope.zero(Elem::UInt);
        let offset_value = scope.zero(Elem::UInt);

        let stride_tensor_dim = scope.create_local(Elem::UInt);
        let stride_value_dim = scope.create_local(Elem::UInt);
        let shape_value_dim = scope.create_local(Elem::UInt);

        let num_elems = scope.create_local(Elem::UInt);
        gpu!(scope, num_elems = cast(1u32));

        gpu!(
            scope,
            range(0u32, Variable::Rank).for_each(|i, scope| {
                let shape_value = scope.create_local(Elem::UInt);
                let stride_tensor = scope.create_local(Elem::UInt);
                let stride_value = scope.create_local(Elem::UInt);

                gpu!(scope, stride_tensor = stride(tensor, i));
                gpu!(scope, stride_value = stride(value, i));
                gpu!(scope, shape_value = shape(value, i));

                let dim_index = scope.create_local(Elem::Bool);
                gpu!(scope, dim_index = i == self.dim);

                gpu!(scope, if(dim_index).then(|scope| {
                    gpu!(scope, shape_value_dim = shape_value);
                    gpu!(scope, stride_tensor_dim = stride_tensor);
                    gpu!(scope, stride_value_dim = stride_value);
                }).else(|scope| {
                    let stride_tmp = scope.create_local(Elem::UInt);
                    let shape_tensor = scope.create_local(Elem::UInt);

                    gpu!(scope, stride_tmp = stride(indices, i));
                    gpu!(scope, shape_tensor = shape(tensor, i));

                    gpu!(scope, num_elems = num_elems * shape_tensor);

                    let offset_local = scope.create_local(Elem::UInt);
                    let offset_local_tensor = scope.create_local(Elem::UInt);
                    let offset_local_value = scope.create_local(Elem::UInt);

                    gpu!(scope, offset_local = id / stride_tmp);

                    gpu!(scope, offset_local_tensor = offset_local % shape_tensor);
                    gpu!(
                        scope,
                        offset_local_tensor = offset_local_tensor * stride_tensor
                    );
                    gpu!(scope, offset_tensor += offset_local_tensor);

                    gpu!(scope, offset_local_value = offset_local % shape_value);
                    gpu!(
                        scope,
                        offset_local_value = offset_local_value * stride_value
                    );
                    gpu!(scope, offset_value += offset_local_value);
                }));
            })
        );

        let should_stop = scope.create_local(Elem::Bool);
        gpu!(scope, should_stop = id >= num_elems);
        gpu!(scope, if(should_stop).then(|scope| {
            scope.register(Branch::Return);
        }));

        gpu!(
            scope,
            range(0u32, shape_value_dim).for_each(|i, scope| {
                let index = scope.create_local(Elem::UInt);
                let index_tensor = scope.create_local(Elem::UInt);
                let index_value = scope.create_local(Elem::UInt);

                let result_tensor = scope.create_local(tensor.item());
                let result_value = scope.create_local(value.item());
                let result = scope.create_local(tensor.item());

                gpu!(scope, index = indices[i]);

                gpu!(scope, index_tensor = index * stride_tensor_dim);
                gpu!(scope, index_tensor += offset_tensor);

                gpu!(scope, index_value = i * stride_value_dim);
                gpu!(scope, index_value += offset_value);

                gpu!(scope, result_tensor = tensor[index_tensor]);
                gpu!(scope, result_value = value[index_value]);
                gpu!(scope, result = result_value + result_tensor);

                gpu!(scope, tensor[index_tensor] = result);
            })
        );
    }