code: def _linear_regression_rows_nd(
      y, x, covariates, block_size=16, weights=None, pass_through=()
  ) -> hail.Table:
      mt = matrix_table_source("linear_regression_rows_nd/x", x)
      raise_unless_entry_indexed("linear_regression_rows_nd/x", x)
  
      y_is_list = isinstance(y, list)
      if y_is_list and len(y) == 0:
          raise ValueError(
              "'linear_regression_rows_nd': found no values for 'y'"
          )
      is_chained = y_is_list and isinstance(y[0], list)
  
      if is_chained and any(len(lst) == 0 for lst in y):
          raise ValueError(
              "'linear_regression_rows': found empty inner list for 'y'"
          )
  
      y = [
          raise_unless_column_indexed("linear_regression_rows_nd/y", y)
          or ys
          for ys in wrap_to_list(y)
          for y in (ys if is_chained else [ys])
      ]
  
      if weights is not None:
          if y_is_list and is_chained and not isinstance(weights, list):
              raise ValueError(
                  "When y is a list of lists, weights should be a list."
              )
          elif (
              y_is_list and not is_chained and isinstance(weights, list)
          ):
              raise ValueError(
                  "When y is a single list, weights should be a single expression."
              )
          elif not y_is_list and isinstance(weights, list):
              raise ValueError(
                  "When y is a single expression, weights should be a single expression."
              )
  
      weights = wrap_to_list(weights) if weights is not None else None
  
      for e in itertools.chain.from_iterable(y) if is_chained else y:
          analyze("linear_regression_rows_nd/y", e, mt._col_indices)
  
      for e in covariates:
          analyze(
              "linear_regression_rows_nd/covariates", e, mt._col_indices
          )
  
      _warn_if_no_intercept("linear_regression_rows_nd", covariates)
  
      x_field_name = Env.get_uid()
      if is_chained:
          y_field_name_groups = [
              [f"__y_{i}_{j}" for j in range(len(y[i]))]
              for i in range(len(y))
          ]
          y_dict = dict(
              zip(
                  itertools.chain.from_iterable(y_field_name_groups),
                  itertools.chain.from_iterable(y),
              )
          )
          if weights is not None and len(weights) != len(y):
              raise ValueError(
                  "Must specify same number of weights as groups of phenotypes"
              )
      else:
          y_field_name_groups = list(f"__y_{i}" for i in range(len(y)))
          y_dict = dict(zip(y_field_name_groups, y))
          # Wrapping in a list since the code is written for the more general chained case.
          y_field_name_groups = [y_field_name_groups]
          if weights is not None and len(weights) != 1:
              raise ValueError(
                  "Must specify same number of weights as groups of phenotypes"
              )
  
      cov_field_names = list(
          f"__cov{i}" for i in range(len(covariates))
      )
      weight_field_names = (
          list(f"__weight_for_group_{i}" for i in range(len(weights)))
          if weights is not None
          else None
      )
      weight_dict = (
          dict(zip(weight_field_names, weights))
          if weights is not None
          else {}
      )
  
      row_field_names = _get_regression_row_fields(
          mt, pass_through, "linear_regression_rows_nd"
      )
  
      # FIXME: selecting an existing entry field should be emitted as a SelectFields
      mt = mt._select_all(
          col_exprs=dict(
              **y_dict,
              **weight_dict,
              **dict(zip(cov_field_names, covariates)),
          ),
          row_exprs=row_field_names,
          col_key=[],
          entry_exprs={x_field_name: x},
      )
  
      entries_field_name = "ent"
      sample_field_name = "by_sample"
  
      num_y_lists = len(y_field_name_groups)
  
      # Given a hail array, get the mean of the nonmissing entries and
      # return new array where the missing entries are the mean.
      def mean_impute(hl_array):
          non_missing_mean = hl.mean(hl_array, filter_missing=True)
          return hl_array.map(
              lambda entry: hl.if_else(
                  hl.is_defined(entry), entry, non_missing_mean
              )
          )
  
      def select_array_indices(hl_array, indices):
          return indices.map(lambda i: hl_array[i])
  
      def dot_rows_with_themselves(matrix):
          return (matrix * matrix).sum(1)
  
      def no_missing(hail_array):
          return hail_array.all(lambda element: hl.is_defined(element))
  
      ht_local = mt._localize_entries(
          entries_field_name, sample_field_name
      )
  
      ht = ht_local.transmute(
          **{
              entries_field_name: ht_local[entries_field_name][
                  x_field_name
              ]
          }
      )
  
      def setup_globals(ht):
          # cov_arrays is per sample, then per cov.
          if covariates:
              ht = ht.annotate_globals(
                  cov_arrays=ht[sample_field_name].map(
                      lambda sample_struct: [
                          sample_struct[cov_name]
                          for cov_name in cov_field_names
                      ]
                  )
              )
          else:
              ht = ht.annotate_globals(
                  cov_arrays=ht[sample_field_name].map(
                      lambda sample_struct: hl.empty_array(hl.tfloat64)
                  )
              )
  
          y_arrays_per_group = [
              ht[sample_field_name].map(
                  lambda sample_struct: [
                      sample_struct[y_name]
                      for y_name in one_y_field_name_set
                  ]
              )
              for one_y_field_name_set in y_field_name_groups
          ]
  
          if weight_field_names:
              weight_arrays = ht[sample_field_name].map(
                  lambda sample_struct: [
                      sample_struct[weight_name]
                      for weight_name in weight_field_names
                  ]
              )
          else:
              weight_arrays = ht[sample_field_name].map(
                  lambda sample_struct: hl.empty_array(hl.tfloat64)
              )
  
          ht = ht.annotate_globals(
              y_arrays_per_group=y_arrays_per_group,
              weight_arrays=weight_arrays,
          )
          ht = ht.annotate_globals(
              all_covs_defined=ht.cov_arrays.map(
                  lambda sample_covs: no_missing(sample_covs)
              )
          )
  
          def get_kept_samples(group_idx, sample_ys):
              # sample_ys is an array of samples, with each element being an array of the y_values
              return (
                  hl.enumerate(sample_ys)
                  .filter(
                      lambda idx_and_y_values: ht.all_covs_defined[
                          idx_and_y_values[0]
                      ]
                      & no_missing(idx_and_y_values[1])
                      & (
                          hl.is_defined(
                              ht.weight_arrays[idx_and_y_values[0]][
                                  group_idx
                              ]
                          )
                          if weights
                          else True
                      )
                  )
                  .map(lambda idx_and_y_values: idx_and_y_values[0])
              )
  
          ht = ht.annotate_globals(
              kept_samples=hl.enumerate(ht.y_arrays_per_group).starmap(
                  get_kept_samples
              )
          )
          ht = ht.annotate_globals(
              y_nds=hl.zip(
                  ht.kept_samples, ht.y_arrays_per_group
              ).starmap(
                  lambda sample_indices, y_arrays: hl.nd.array(
                      sample_indices.map(lambda idx: y_arrays[idx])
                  )
              )
          )
          ht = ht.annotate_globals(
              cov_nds=ht.kept_samples.map(
                  lambda group: hl.nd.array(
                      group.map(lambda idx: ht.cov_arrays[idx])
                  )
              )
          )
  
          if weights is None:
              ht = ht.annotate_globals(
                  sqrt_weights=hl.missing(
                      hl.tarray(hl.tndarray(hl.tfloat64, 2))
                  )
              )
              ht = ht.annotate_globals(scaled_y_nds=ht.y_nds)
              ht = ht.annotate_globals(scaled_cov_nds=ht.cov_nds)
          else:
              ht = ht.annotate_globals(
                  weight_nds=hl.enumerate(ht.kept_samples).starmap(
                      lambda group_idx, group_sample_indices: hl.nd.array(
                          group_sample_indices.map(
                              lambda group_sample_idx: ht.weight_arrays[
                                  group_sample_idx
                              ][group_idx]
                          )
                      )
                  )
              )
              ht = ht.annotate_globals(
                  sqrt_weights=ht.weight_nds.map(
                      lambda weight_nd: weight_nd.map(
                          lambda e: hl.sqrt(e)
                      )
                  )
              )
              ht = ht.annotate_globals(
                  scaled_y_nds=hl.zip(
                      ht.y_nds, ht.sqrt_weights
                  ).starmap(
                      lambda y, sqrt_weight: y
                      * sqrt_weight.reshape(-1, 1)
                  )
              )
              ht = ht.annotate_globals(
                  scaled_cov_nds=hl.zip(
                      ht.cov_nds, ht.sqrt_weights
                  ).starmap(
                      lambda cov, sqrt_weight: cov
                      * sqrt_weight.reshape(-1, 1)
                  )
              )
  
          k = builtins.len(covariates)
          ht = ht.annotate_globals(
              ns=ht.kept_samples.map(
                  lambda one_sample_set: hl.len(one_sample_set)
              )
          )
  
          def log_message(i):
              if is_chained:
                  return (
                      "linear regression_rows["
                      + hl.str(i)
                      + "] running on "
                      + hl.str(ht.ns[i])
                      + " samples for "
                      + hl.str(ht.scaled_y_nds[i].shape[1])
                      + f" response variables y, with input variables x, and {len(covariates)} additional covariates..."
                  )
              else:
                  return (
                      "linear_regression_rows running on "
                      + hl.str(ht.ns[0])
                      + " samples for "
                      + hl.str(ht.scaled_y_nds[i].shape[1])
                      + f" response variables y, with input variables x, and {len(covariates)} additional covariates..."
                  )
  
          ht = ht.annotate_globals(
              ns=hl.range(num_y_lists).map(
                  lambda i: hl._console_log(log_message(i), ht.ns[i])
              )
          )
          ht = ht.annotate_globals(
              cov_Qts=hl.if_else(
                  k > 0,
                  ht.scaled_cov_nds.map(
                      lambda one_cov_nd: hl.nd.qr(one_cov_nd)[0].T
                  ),
                  ht.ns.map(lambda n: hl.nd.zeros((0, n))),
              )
          )
          ht = ht.annotate_globals(
              Qtys=hl.zip(ht.cov_Qts, ht.scaled_y_nds).starmap(
                  lambda cov_qt, y: cov_qt @ y
              )
          )
  
          return ht.select_globals(
              kept_samples=ht.kept_samples,
              __scaled_y_nds=ht.scaled_y_nds,
              __sqrt_weight_nds=ht.sqrt_weights,
              ns=ht.ns,
              ds=ht.ns.map(lambda n: n - k - 1),
              __cov_Qts=ht.cov_Qts,
              __Qtys=ht.Qtys,
              __yyps=hl.range(num_y_lists).map(
                  lambda i: dot_rows_with_themselves(
                      ht.scaled_y_nds[i].T
                  )
                  - dot_rows_with_themselves(ht.Qtys[i].T)
              ),
          )
  
      ht = setup_globals(ht)
  
      def process_block(block):
          rows_in_block = hl.len(block)
  
          # Processes one block group based on given idx. Returns a single struct.
          def process_y_group(idx):
              if weights is not None:
                  X = (
                      hl.nd.array(
                          block[entries_field_name].map(
                              lambda row: mean_impute(
                                  select_array_indices(
                                      row, ht.kept_samples[idx]
                                  )
                              )
                          )
                      )
                      * ht.__sqrt_weight_nds[idx]
                  ).T
              else:
                  X = hl.nd.array(
                      block[entries_field_name].map(
                          lambda row: mean_impute(
                              select_array_indices(
                                  row, ht.kept_samples[idx]
                              )
                          )
                      )
                  ).T
              n = ht.ns[idx]
              sum_x = X.sum(0)
              Qtx = ht.__cov_Qts[idx] @ X
              ytx = ht.__scaled_y_nds[idx].T @ X
              xyp = ytx - (ht.__Qtys[idx].T @ Qtx)
              xxpRec = (
                  dot_rows_with_themselves(X.T)
                  - dot_rows_with_themselves(Qtx.T)
              ).map(lambda entry: 1 / entry)
              b = xyp * xxpRec
              se = (
                  (1.0 / ht.ds[idx])
                  * (
                      ht.__yyps[idx].reshape((-1, 1))
                      @ xxpRec.reshape((1, -1))
                      - (b * b)
                  )
              ).map(lambda entry: hl.sqrt(entry))
              t = b / se
              return hl.rbind(
                  t,
                  lambda t: hl.rbind(
                      ht.ds[idx],
                      lambda d: hl.rbind(
                          t.map(
                              lambda entry: 2
                              * hl.expr.functions.pT(
                                  -hl.abs(entry), d, True, False
                              )
                          ),
                          lambda p: hl.struct(
                              n=hl.range(rows_in_block).map(
                                  lambda i: n
                              ),
                              sum_x=sum_x._data_array(),
                              y_transpose_x=ytx.T._data_array(),
                              beta=b.T._data_array(),
                              standard_error=se.T._data_array(),
                              t_stat=t.T._data_array(),
                              p_value=p.T._data_array(),
                          ),
                      ),
                  ),
              )
  
          per_y_list = hl.range(num_y_lists).map(
              lambda i: process_y_group(i)
          )
  
          key_field_names = [key_field for key_field in ht.key]
  
          def build_row(row_idx):
              # For every field we care about, map across all y's, getting the row_idxth one from each.
              idxth_keys = {
                  field_name: block[field_name][row_idx]
                  for field_name in key_field_names
              }
              computed_row_field_names = [
                  "n",
                  "sum_x",
                  "y_transpose_x",
                  "beta",
                  "standard_error",
                  "t_stat",
                  "p_value",
              ]
              computed_row_fields = {
                  field_name: per_y_list.map(
                      lambda one_y: one_y[field_name][row_idx]
                  )
                  for field_name in computed_row_field_names
              }
              pass_through_rows = {
                  field_name: block[field_name][row_idx]
                  for field_name in row_field_names
              }
  
              if not is_chained:
                  computed_row_fields = {
                      key: value[0]
                      for key, value in computed_row_fields.items()
                  }
  
              return hl.struct(
                  **{
                      **idxth_keys,
                      **computed_row_fields,
                      **pass_through_rows,
                  }
              )
  
          new_rows = hl.range(rows_in_block).map(build_row)
  
          return new_rows
  
      def process_partition(part):
          grouped = part.grouped(block_size)
          return grouped.flatmap(
              lambda block: process_block(block)._to_stream()
          )
  
      res = ht._map_partitions(process_partition)
  
      if not y_is_list:
          fields = [
              "y_transpose_x",
              "beta",
              "standard_error",
              "t_stat",
              "p_value",
          ]
          res = res.annotate(**{f: res[f][0] for f in fields})
  
      res = res.select_globals()
  
      temp_file_name = hl.utils.new_temp_file(
          "_linear_regression_rows_nd", "result"
      )
      res = res.checkpoint(temp_file_name)
  
      return res
  
imports: import builtins
  import itertools
  import math
  from typing import Dict, Callable, Optional, Union, Tuple, List
  import hail
  import hail
  import hail.expr.aggregators
  from hail import ir
  from hail.expr import Expression, ExpressionException, expr_float64, expr_call, expr_any, expr_numeric, expr_locus, analyze, raise_unless_entry_indexed, raise_unless_row_indexed, matrix_table_source, table_source, raise_unless_column_indexed
  from hail.expr.types import tbool, tarray, tfloat64, tint32
  from hail.genetics.reference_genome import reference_genome_type
  from hail.linalg import BlockMatrix
  from hail.matrixtable import MatrixTable
  from hail.methods.misc import require_biallelic, require_row_key_variant
  from hail.stats import LinearMixedModel
  from hail.table import Table
  from hail.typecheck import typecheck, nullable, numeric, oneof, sized_tupleof, sequenceof, enumeration, anytype
  from hail.utils import wrap_to_list, new_temp_file, FatalError
  from hail.utils.java import Env, info, warning
  from .None import pca
  from .None import relatedness
  from ..backend.spark_backend import SparkBackend
label_desc: Use image node_id 4 for calling a global function with example usage: _linear_regression_rows_nd(y, x, covariates, block_size, weights, pass_through) and returns: res, hl_array, indices, unknown, hail_array, ht, new_rows, grouped, hl, hl, hl, unknown, unknown