# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
#   """Expand unbatched `x` to the same batch size as `y`."""
#   batch_size = y.shape[0]
#   return jnp.tile(
#       jnp.expand_dims(x, axis=0), [batch_size] + [1 for _ in x.shape])
# 
# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
# def np_load(path: str, restore_dtype: Optional[DType] = None) -> Array:
# def t5x_load(checkpoint_path: str,
#              variable_path: str,
#              restore_dtype: Optional[DType] = None) -> Array:
# def from_array(prompt: Array) -> Initializer:
#   def initialize_from_file(unused_prng_key: Array,
#                            shape: Sequence[int]) -> Array:
# def from_sample_of_embeddings(
#     embeddings: Array,
#     population_size: Optional[int] = None,
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedding_sample(rng: Array,
#                                        shape: Sequence[int]) -> Array:
# def from_embedded_list(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     texts: Sequence[str],
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_list(rng: Array, shape: Sequence[int]) -> Array:
# def from_embedded_string(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     text: str,
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_string(rng: Array,
#                                       shape: Sequence[int]) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


from typing import Callable, Optional, Any, Iterable, Tuple
from flax import linen as nn
from flax.linen import partitioning
from prompt_tuning.prompts import expand_to_batch
from prompt_tuning.prompts import Initializer
from flaxformer.components import embedding
from flaxformer.types import Array
import jax.numpy as jnp
  task_features = shared_features
  if self.task_features is not None:
    task_features = self.task_features
  kernel = shared_features * task_features
  bias = task_features
  return kernel + bias


# TODO: Update these classes to use actual subclasses with
# overridden methods once that problem is fixed.
class MultiTaskPrompt(nn.Module):
  """A module that produces a, combined, shared and tasks specific prompt.

  Attributes:
    length: The length of the shared prompt, P_s.
    num_tasks: The number of tasks that are including in our multi-task setup.
    to_task_idx: A function that maps the first token of each element in the
      batch to a task index starting from zero.
    shared_init: An initializer for the shared prompt variables.
    task_init: An initializer for the task prompts variables.
    get_shared_features: A function (pretending to be a method) that calculates
      the number of features in the shared prompt.
    get_task_size: A function (pretending to be a method) that calculates the
      number of values in a single (flattened) task specific prompt.
    combine_shared_and_tasks: A function that combines the shared and
      (flattened) task specific prompts.
  """
  length: int
  num_tasks: int
  to_task_idx: Callable[[Array], Array] = identity

# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
#   """Expand unbatched `x` to the same batch size as `y`."""
#   batch_size = y.shape[0]
#   return jnp.tile(
#       jnp.expand_dims(x, axis=0), [batch_size] + [1 for _ in x.shape])
# 
# Path: prompt_tuning/prompts.py
# def expand_to_batch(x: Array, y: Array):
# def np_load(path: str, restore_dtype: Optional[DType] = None) -> Array:
# def t5x_load(checkpoint_path: str,
#              variable_path: str,
#              restore_dtype: Optional[DType] = None) -> Array:
# def from_array(prompt: Array) -> Initializer:
#   def initialize_from_file(unused_prng_key: Array,
#                            shape: Sequence[int]) -> Array:
# def from_sample_of_embeddings(
#     embeddings: Array,
#     population_size: Optional[int] = None,
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedding_sample(rng: Array,
#                                        shape: Sequence[int]) -> Array:
# def from_embedded_list(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     texts: Sequence[str],
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_list(rng: Array, shape: Sequence[int]) -> Array:
# def from_embedded_string(
#     embeddings: Array,
#     vocab: seqio.SentencePieceVocabulary,
#     text: str,
#     initializer: Optional[Callable[[Array, Sequence[int]], Array]] = None
# ) -> Callable[[Array, Sequence[int]], Array]:
#   def initialize_from_embedded_string(rng: Array,
#                                       shape: Sequence[int]) -> Array:
#   def __call__(self, x, x_embed):
# class Prompt(nn.Module):


# Path: prompt_tuning/extended/multitask_prompts.py
from typing import Callable, Optional, Any, Iterable, Tuple
from flax import linen as nn
from flax.linen import partitioning
from prompt_tuning.prompts import expand_to_batch
from prompt_tuning.prompts import Initializer
from flaxformer.components import embedding
from flaxformer.types import Array
import jax.numpy as jnp

  task_features = shared_features
  if self.task_features is not None:
    task_features = self.task_features
  kernel = shared_features * task_features
  bias = task_features
  return kernel + bias


# TODO: Update these classes to use actual subclasses with
# overridden methods once that problem is fixed.
class MultiTaskPrompt(nn.Module):
  """A module that produces a, combined, shared and tasks specific prompt.

  Attributes:
    length: The length of the shared prompt, P_s.
    num_tasks: The number of tasks that are including in our multi-task setup.
    to_task_idx: A function that maps the first token of each element in the
      batch to a task index starting from zero.
    shared_init: An initializer for the shared prompt variables.
    task_init: An initializer for the task prompts variables.
    get_shared_features: A function (pretending to be a method) that calculates
      the number of features in the shared prompt.
    get_task_size: A function (pretending to be a method) that calculates the
      number of values in a single (flattened) task specific prompt.
    combine_shared_and_tasks: A function that combines the shared and
      (flattened) task specific prompts.
  """
  length: int
  num_tasks: int
  to_task_idx: Callable[[Array], Array] = identity

  shared_init: Initializer = nn.initializers.uniform()