text:
"""Compile stages to pipeline instructions."""
from collections import namedtuple, defaultdict
from dataclasses import dataclass
import enum
import logging
from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Union

from jax._src.tree_util import PyTreeDef, tree_unflatten
from jax.core import Var
from jax.interpreters import pxla
from jax.lib import xla_bridge as xb
import numpy as np

from alpa.global_env import global_config
from alpa.device_mesh import DistributedArray, PhysicalDeviceMeshGroup, ReplicatedDistributedArray
from alpa.mesh_executable import next_mesh_executable_uuid
from alpa.pipeline_parallel.computation import XlaShardedPipelineComputation
from alpa.pipeline_parallel.schedules import PipelineSchedule
from alpa.pipeline_parallel.cross_mesh_resharding import (
    CrossMeshCommunicator, SymbolicBroadcastReshardingTask,
    SymbolicReshardingTask)
from alpa.serialization import LoadInfo
from alpa.util import (DisjointDict, OrderedSet, cached_property,
                       get_shard_shape, get_microbatch_sharding_spec,
                       compile_concatenate)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class PipelineInstType(enum.IntEnum):
    """Enum class for pipeline instruction types."""

    # Run an XLA executable
    RUN = 0
    # Run a sending task
    SEND = 1
    # Run a receiving task
    RECV = 2
    # Free tensors
    FREE = 3
    BROADCAST = 4

@dataclass
class PipelineInstruction:
    """Base class for pipeline instructions."""

    opcode: PipelineInstType
    task_uuid: Optional[int]
    input_uuids: Optional[np.ndarray]
    output_uuids: Optional[np.ndarray]
    opaques: Optional[Dict[str, Any]]
    info: str
    print_uuids: bool = False

    @classmethod
    def Run(cls, task_uuid, input_uuids, output_uuids, kwargs, info=""):  # noqa
        return cls(opcode=PipelineInstType.RUN,
                   task_uuid=task_uuid,
                   input_uuids=input_uuids,
                   output_uuids=output_uuids,
                   opaques={"kwargs": kwargs},
                   info=info)

    @classmethod
    def Send(cls, task_uuid, input_uuids, info=""):  # noqa
        return cls(opcode=PipelineInstType.SEND,
                   task_uuid=task_uuid,
                   input_uuids=input_uuids,
                   output_uuids=None,
                   opaques=None,
                   info=info)

    @classmethod
    def Recv(cls,  # noqa
             task_uuid,
             output_uuids,
             set_empty_buffer,
             allgather_uuid=None,
             info=""):  # noqa
        return cls(opcode=PipelineInstType.RECV,
                   task_uuid=task_uuid,
                   input_uuids=None,
                   output_uuids=output_uuids,
                   opaques={
                       "set_empty_buffer": set_empty_buffer,
                       "allgather_uuid": allgather_uuid
                   },
                   info=info)

    @classmethod
    def Broadcast(cls,  # noqa
                 task_uuid,
                 input_uuids,
                 output_uuids,
                 info="broadcast"):  # noqa
        return cls(opcode=PipelineInstType.BROADCAST,
                   task_uuid=task_uuid,
                   input_uuids=input_uuids,
                   output_uuids=output_uuids,
                   opaques=None,
                   info=info)

    @classmethod
    def Free(cls, input_uuids, info=""):  # noqa
        return cls(opcode=PipelineInstType.FREE,
                   task_uuid=None,
                   input_uuids=input_uuids,
                   output_uuids=None,
                   opaques=None,
                   info=info,
                   print_uuids=False)

    def __str__(self):
        ret = ""
        ret += "Opcode: " + str(self.opcode)[17:] + ", Task uuid: " + str(
            self.task_uuid)
        if self.print_uuids:
            ret += ", input uuids:" + str(self.input_uuids)
            ret += ", output uuids:" + str(self.output_uuids)
        ret += ", Info: " + self.info
        return ret

AllocateZeroWorkerExecutableConfig = namedtuple(
    "AllocateZeroWorkerExecutableConfig",
    ["exec_uuid", "grad_shard_shapes", "grad_shard_dtypes"])
MemZeroWorkerExecutableConfig = namedtuple(
    "MemZeroWorkerExecutableConfig",
    ["exec_uuid", "grad_shard_shapes", "grad_shard_dtypes"])
ConcatWorkerExecutableConfig = namedtuple("ConcatWorkerExecutableConfig",
                                          ["exec_uuid", "hlo_proto"])
PartialGradWorkerExecutableConfig = namedtuple(
    "PartialGradWorkerExecutableConfig",
    ["exec_uuid", "hlo_proto", "strategy_config", "grad_sync_channel_ids"])

ExecutableConfig = Union[AllocateZeroWorkerExecutableConfig,
                         MemZeroWorkerExecutableConfig,
                         PartialGradWorkerExecutableConfig,
                         ConcatWorkerExecutableConfig]

def flatten_uuid_set(container):
    """Convert a nested array to an OrderedSet of elements in the array."""
    output = OrderedSet()
    for e in container:
        if isinstance(e, (np.ndarray, list)):
            output.update(flatten_uuid_set(e))
        else:
            output.add(e)
    return output

class PipelineInstEmitterHelper:

    def __init__(self, global_invars, grad_dummy_invars, is_batch,
                 schedule: PipelineSchedule):
        self.global_invars = global_invars
        self.global_batch_invars = OrderedSet(
            v for v, b in zip(global_invars, is_batch) if b)
        self.grad_dummy_invars = grad_dummy_invars
        self.schedule = schedule
        # Dict[var_key -> Dict[mesh_idx -> np.array of uuids]]
        # The shape of the numpy array is [num_hosts, num_devices_per_host]
        self.env = {}

    def _get_var_key(self, var, batch_idx):
        if var in self.global_invars and var not in self.global_batch_invars:
            key = (var, 0)
        elif (var in self.grad_dummy_invars and
              batch_idx != self.schedule.first_backward_batch_index):
            key = (self.grad_dummy_invars[var],
                   self.schedule.previous_backward_batch_index(batch_idx))
        else:
            key = (var, batch_idx)
        return key

    def get_var_with_accumulate(self, var, batch_idx):
        if (var in self.grad_dummy_invars and
                batch_idx != self.schedule.first_backward_batch_index):
            return self.grad_dummy_invars[var]
        else:
            return var

    def get_var_mesh_uuids(self, var, batch_idx, mesh_idx) -> np.ndarray:
        key = self._get_var_key(var, batch_idx)
        return self.env[key][mesh_idx]

    def get_var_meshes(self, var, batch_idx) -> Dict[int, np.ndarray]:
        key = self._get_var_key(var, batch_idx)
        return self.env.setdefault(key, {})

    def set_var_mesh_uuids(self, var, batch_idx, mesh_idx, uuids) -> None:
        key = self._get_var_key(var, batch_idx)
        self.env.setdefault(key, {})[mesh_idx] = uuids

    def var_at(self, var, batch_idx, mesh_idx) -> bool:
        key = self._get_var_key(var, batch_idx)
        return mesh_idx in self.env.setdefault(key, {})

@dataclass
class PipeshardInputConfig:
    input_local_uuid_lists: Dict[Any, Sequence[np.ndarray]]
    donate_invars: Sequence[Sequence[bool]]
    mesh_arg_indices: Sequence[Sequence[int]]
    input_shard_indices: Sequence[Sequence[Any]]
    delete_after_shard: Sequence[Sequence[bool]]
    batch_invars: Sequence[Sequence[bool]]

# TODO(yonghao): use worker_idx as the dict's key
@dataclass
class PipeshardConfig:
    instruction_lists: Dict[Any, Sequence[PipelineInstruction]]
    # Executable configs
    executable_configs: Dict[Any, Sequence[ExecutableConfig]]
    executable_uuids: Sequence[int]
    # Input configs
    input_config: PipeshardInputConfig
    grad_uuids: Sequence[np.ndarray]
    reduced_var_uuid_lists: Dict[Any, np.ndarray]
    device_str_groups: Sequence[Sequence[OrderedSet]]
    # Output configs
    output_local_uuid_list: Sequence[Sequence]
    # Others
    outs_handler: Callable
    load_info: LoadInfo
    resharding_task_iter: Iterator

    @property
    def input_local_uuid_lists(self):
        return self.input_config.input_local_uuid_lists

    @property
    def donate_invars(self):
        return self.input_config.donate_invars

    @property
    def mesh_arg_indices(self):
        return self.input_config.mesh_arg_indices

    @property
    def input_shard_indices(self):
        return self.input_config.input_shard_indices

    @property
    def delete_after_shard(self):
        return self.input_config.delete_after_shard

    @property
    def batch_invars(self):
        return self.input_config.batch_invars

class PipelineInstEmitter:

    def __init__(self, *, stages: Sequence[XlaShardedPipelineComputation],
                 global_invars: Sequence[Var], grad_dummy_invars: Sequence[Var],
                 concat_vars_mapping: Dict[Var, Var],
                 global_outvars: Sequence[Var],
                 mesh_group: PhysicalDeviceMeshGroup,
                 schedule: PipelineSchedule, is_batch: Sequence[bool],
                 num_batch: int, in_tree: PyTreeDef):
        ##### Input arguments #####
        self.stages = stages
        self.global_invars = global_invars
        self.grad_dummy_invars = grad_dummy_invars
        self.concat_vars_mapping = concat_vars_mapping
        self.global_outvars = global_outvars
        self.mesh_group = mesh_group
        self.schedule = schedule
        self.is_batch = is_batch
        self.num_batch = num_batch
        self.in_tree = in_tree

        ##### Internal states #####
        self.uuid_counter = 0  # counter for local buffer uuid
        self.env = PipelineInstEmitterHelper(global_invars, grad_dummy_invars,
                                             is_batch, schedule)

        # for cross mesh communications
        self.num_mesh = len(mesh_group)

        ##### For handling outputs of the executable ####
        # List[arg_idx -> List[mesh_idx -> int]]
        self.mesh_output_indices = []
        # List[mesh_idx -> List[arg_idx -> sharding_spec]]
        self.output_spec_list = [[] for _ in range(len(mesh_group))]

        ##### For cross mesh communication ####
        self._precompile_sharding_specs()
        self._communicator = CrossMeshCommunicator(self.stages, self.schedule)
        self.num_mesh = len(mesh_group)
        self._resharding_tasks = [
            [{} for _ in range(self.num_mesh)] for _ in range(self.num_mesh)
        ]

    def _get_next_uuids(self, num) -> np.ndarray:
        """Get the next uuids as a numpy array of uuids."""
        ret = np.arange(start=self.uuid_counter,
                        stop=self.uuid_counter + num,
                        dtype=np.int64)
        self.uuid_counter += num
        return ret

    def _precompile_sharding_specs(self):
        """Run spmd partitioner pass for each stage to get sharding specs."""
        for stage_idx, stage in enumerate(self.stages):
            mesh_indices = list(self.schedule.stage_placement(stage_idx))
            assert len(mesh_indices) == 1
            stage.get_spmd_partitioned()

    def _compile_resharding_tasks(self):
        """Create and compile all resharding (send/recv/allgather) tasks."""
        for (src_mesh_idx, dst_mesh_idx,
             var_spec_map) in self._communicator.task_spec_iter():
            for var, spec in var_spec_map.items():
                cg = self.mesh_group.collective_groups[src_mesh_idx][
                    dst_mesh_idx]
                src_mesh = self.mesh_group[src_mesh_idx]
                dst_mesh = self.mesh_group[dst_mesh_idx]
                # TODO(yonghao): delay put_resharding_XXXX_task until pipeshard executable
                if global_config.resharding_mode == "send_recv":
                    self._resharding_tasks[src_mesh_idx][dst_mesh_idx][
                        var] = SymbolicReshardingTask(spec, cg, src_mesh,
                                                      dst_mesh)
                else:
                    self._resharding_tasks[src_mesh_idx][dst_mesh_idx][
                        var] = SymbolicBroadcastReshardingTask(
                            spec, cg, src_mesh, dst_mesh)

    def _establish_nccl_groups(self):
        """
        Identify NCCL groups based on resharding specs but do not instantiate them.

        We establish one collective group between two physical meshes, covering all the devices in
        these two meshes that require NCCL communication.

        Returns:
            device_str_groups (List[List[set]]): a num_mesh x num_mesh matrix. Only entries at
                device_str_groups[i][j] (i < j) are filled, entries with i > j are None, because
                (spec[i][j], spec[j][i]) will share collective groups.
        """
        device_str_groups = [[OrderedSet()
                              for _ in range(self.num_mesh)]
                             for _ in range(self.num_mesh)]
        # Merge (i, j) and (j, i)
        for i, j, var_spec_map in self._communicator.task_spec_iter():
            participants = OrderedSet()
            for _, spec in var_spec_map.items():  # for each var
                participants = participants | spec.get_participant_device_strs()
            if i <= j:
                device_str_groups[i][j] = device_str_groups[i][j] | participants
            else:
                device_str_groups[j][i] = device_str_groups[j][i] | participants

        # construct groups
        for i in range(self.num_mesh):
            for j in range(self.num_mesh):
                if i >= j:
                    assert not device_str_groups[i][j]
                    continue
                if not device_str_groups[i][j]:
                    continue
                self.mesh_group.establish_nccl_group(i, j, instantiate=False)
        return device_str_groups

    def compile(self):
        """Compile pipeline instructions and executables for workers."""
        num_mesh = len(self.mesh_group)
        device_str_groups = self._establish_nccl_groups()

        self._compile_resharding_tasks()
        # Compile forward, backward and apply_grad computations
        (executable_uuids,
         executable_config_lists) = self._compile_computation_executables()

        # Compile gradient buffer allocations
        grad_uuids, instruction_lists = self._compile_grad_buffer_allocations(
            executable_config_lists)

        # Split input into micro batches
        global_batch_invar_set = OrderedSet([
            var for var, batch in zip(self.global_invars, self.is_batch)
            if batch
        ])
        (input_config, input_shard_specs
         ) = self._compile_split_input_to_microbatches(global_batch_invar_set)

        # Simulate the pipeline schedule and generate instructions
        donation_mapping = [DisjointDict() for _ in range(num_mesh)]
        worker_to_idx = {}
        for mesh_idx, mesh in enumerate(self.mesh_group):
            for worker_idx, worker in enumerate(mesh.workers):
                worker_to_idx[worker] = (mesh_idx, worker_idx)

        for _, sched in enumerate(self.schedule.schedules):
            self._compile_exec_one_tick(sched, donation_mapping,
                                        instruction_lists, executable_uuids,
                                        executable_config_lists)

        # Compile concate
        self._compile_concate(instruction_lists, executable_config_lists)

        # Compile information for outputs
        output_local_uuid_list = self._compile_collect_outputs()

        # Insert buffer free instructions
        reduced_var_uuid_lists = {}
        for worker in instruction_lists:
            used_outside = flatten_uuid_set(output_local_uuid_list[worker])
            mesh_idx, worker_idx = worker_to_idx[worker]
            reduced_var_uuids = grad_uuids[mesh_idx]
            if len(reduced_var_uuids) > 0:
                reduced_var_uuids = reduced_var_uuids[worker_idx]
            reduced_var_uuids = np.array([[
                donation_mapping[mesh_idx].recursive_lookup(uuid)
                for uuid in uuids
            ] for uuids in reduced_var_uuids])
            donated = set(donation_mapping[mesh_idx].keys())
            used_outside.update(flatten_uuid_set(reduced_var_uuids))
            reduced_var_uuid_lists[worker] = reduced_var_uuids
            instruction_lists[worker] = self._compile_free(
                worker, used_outside, donated, instruction_lists)

        outs_handler = self._get_outs_handler()

        load_info = self._compile_load_info(input_config.mesh_arg_indices,
                                            input_shard_specs)
        return PipeshardConfig(
            instruction_lists=instruction_lists,
            # Executable configs
            executable_configs=executable_config_lists,
            executable_uuids=executable_uuids,
            # Input configs
            input_config=input_config,
            grad_uuids=grad_uuids,
            reduced_var_uuid_lists=reduced_var_uuid_lists,
            device_str_groups=device_str_groups,
            # Output configs
            output_local_uuid_list=output_local_uuid_list,
            # Others
            outs_handler=outs_handler,
            load_info=load_info,
            resharding_task_iter=self.resharding_task_iter())

    def resharding_task_iter(self):
        """An iterator over all resharding tasks."""
        for src_idx in range(self.num_mesh):
            for dst_idx in range(self.num_mesh):
                tasks = self._resharding_tasks[src_idx][dst_idx].values()
                for task in tasks:
                    yield task

    def _compile_get_vars_from_mesh(self, invars, dst_specs, mesh_idx,
                                    batch_idx, instruction_lists,
                                    executable_config_lists):
        if len(invars) == 0:
            return
        # TODO(yonghao): only compile alloc once, use multiple times
        output_uuids = self._compile_alloc(invars, dst_specs, mesh_idx,
                                           batch_idx, False, instruction_lists,
                                           executable_config_lists)
        # shape: (args, num_hosts, num_devices_per_host)
        transposed = output_uuids.transpose([1, 0, 2])
        recv_uuid_list = transposed

        for invar, recv_uuids in zip(invars, recv_uuid_list):
            var_key = self.env.get_var_with_accumulate(invar, batch_idx)
            src_idx, src_uuids = list(
                self.env.get_var_meshes(invar, batch_idx).items())[0]
            resharding_task = self._resharding_tasks[src_idx][mesh_idx][var_key]
            if global_config.resharding_mode == "send_recv":
                self._compile_resharding_task(self.mesh_group[src_idx],
                                              self.mesh_group[mesh_idx],
                                              src_uuids, resharding_task,
                                              recv_uuids, instruction_lists)
            else:
                self._compile_broadcast_resharding_task(
                    self.mesh_group[src_idx], self.mesh_group[mesh_idx],
                    src_uuids, resharding_task, recv_uuids, instruction_lists)

    def _compile_exec_one_tick(self, sched, donation_mapping, instruction_lists,
                               executable_uuids, executable_config_lists):
        worker_tmp_instructions = {}
        for mesh in self.mesh_group:
            for worker in mesh.workers:
                worker_tmp_instructions[worker] = []

        for mesh_idx, task in enumerate(sched):
            if not task:
                continue
            physical_mesh = self.mesh_group[mesh_idx]
            num_devices_per_host = physical_mesh.num_devices_per_host
            batch_idx, stage_idx = task
            stage = self.stages[stage_idx]
            # shard_args for intermediates
            to_reshard_vars = []
            reshard_sharding_specs = []
            for invar, spec in zip(stage.invars, stage.input_sharding_specs):
                if self.env.var_at(invar, batch_idx, mesh_idx):
                    # have a copy at the current mesh
                    continue
                if len(self.env.get_var_meshes(invar, batch_idx)) > 1:
                    raise NotImplementedError(
                        "Not support resharding replicated")
                to_reshard_vars.append(invar)
                reshard_sharding_specs.append(spec)
            self._compile_get_vars_from_mesh(to_reshard_vars,
                                             reshard_sharding_specs, mesh_idx,
                                             batch_idx, instruction_lists,
                                             executable_config_lists)

            # execute
            # allocate uuids for buffers created by RUN
            for outvar in stage.outvars:
                # get uuids of this outvar

                output_uuids = self._get_next_uuids(
                    physical_mesh.num_devices).reshape(-1, num_devices_per_host)
                self.env.set_var_mesh_uuids(outvar, batch_idx, mesh_idx,
                                            output_uuids)

            exec_uuid = executable_uuids[stage_idx]
            donated_invars = self.stages[stage_idx].donated_invars

            input_uuids = np.zeros((len(stage.invars), *physical_mesh.shape),
                                   dtype=np.int64)
            output_uuids = np.zeros((len(stage.outvars), *physical_mesh.shape),
                                    dtype=np.int64)
            for idx, invar in enumerate(stage.invars):
                input_uuids[idx] = self.env.get_var_mesh_uuids(
                    invar, batch_idx, mesh_idx)
            for idx, outvar in enumerate(stage.outvars):
                output_uuids[idx] = self.env.get_var_mesh_uuids(
                    outvar, batch_idx, mesh_idx)
            input_uuids = input_uuids.transpose([1, 0, 2])
            output_uuids = output_uuids.transpose([1, 0, 2])

            for worker_idx, worker in enumerate(physical_mesh.workers):
                # Get input and output uuids. They should be at the mesh
                worker_input_uuids = input_uuids[worker_idx]
                worker_output_uuids = output_uuids[worker_idx]
                for idx in range(len(stage.invars)):
                    if donated_invars[idx]:
                        donation_mapping[mesh_idx].update(
                            worker_input_uuids[idx], worker_output_uuids[idx])

                kwargs = {
                    "skip_grad_sync": self.schedule.should_skip_grad_sync(task),
                    "sync_before": False,
                    "sync_after": False,
                }

                worker_tmp_instructions[worker].append(
                    PipelineInstruction.Run(exec_uuid,
                                            worker_input_uuids,
                                            worker_output_uuids,
                                            kwargs,
                                            info=f"stage {stage_idx}"))

        for worker, worker_instruction in worker_tmp_instructions.items():
            instruction_lists[worker].extend(worker_instruction)

    def _compile_computation_executables(self):
        """Compile executables for forward, backward, and apply_grad compuations."""
        executable_uuids = []  # List[stage_idx -> executable_uuids]
        executable_config_lists = defaultdict(
            list)  # Dict[worker -> List[ExecutableConfig]]

        for stage_idx, stage in enumerate(self.stages):
            exec_uuid = next_mesh_executable_uuid()
            executable_uuids.append(exec_uuid)

            mesh_idx = self.schedule.stage_placement(stage_idx)
            assert len(mesh_idx) == 1
            mesh_idx = list(mesh_idx)[0]
            hlo_module = stage.get_spmd_partitioned()
            hlo_proto = hlo_module.as_serialized_hlo_module_proto()
            exec_config = PartialGradWorkerExecutableConfig(
                exec_uuid, hlo_proto, stage.strategy_config,
                stage.output_acc_grad_indices)
            for worker in self.mesh_group[mesh_idx].workers:
                executable_config_lists[worker].append(exec_config)

        return executable_uuids, executable_config_lists

    def _compile_grad_buffer_allocations(self, executable_config_lists):
        """Compile gradient buffer allocations."""
        num_mesh = len(self.mesh_group)
        mesh_grad_vars = [{} for _ in range(num_mesh)]
        instruction_lists = defaultdict(
            list)  # Dict[worker -> List[PipelineInstruction]]
        # collect gradient accumulation buffers in each mesh
        for stage_idx, stage in enumerate(self.stages):
            mesh_indices = list(self.schedule.stage_placement(stage_idx))
            assert len(mesh_indices) == 1
            mesh_idx = mesh_indices[0]
            grad_var_spec_dict = mesh_grad_vars[mesh_idx]
            input_specs = stage.input_sharding_specs
            for var_idx, invar in enumerate(stage.invars):
                if invar in self.grad_dummy_invars:
                    if invar in grad_var_spec_dict:
                        raise NotImplementedError(
                            f"accumulate {invar} at multiple stages in a mesh")
                    grad_var_spec_dict[invar] = input_specs[var_idx]

        grad_uuids = [[] for _ in range(num_mesh)]
        for mesh_idx in range(num_mesh):
            grad_var_spec_dict = mesh_grad_vars[mesh_idx]
            if len(grad_var_spec_dict):
                grad_vars, grad_sharding_specs = list(
                    zip(*grad_var_spec_dict.items()))

                # TODO(yonghao): Some var has non-gradient intermediate states
                # that need accumulation. for these vars, we need to record its
                # first mb index when accum will take place.
                grad_uuids[mesh_idx] = self._compile_alloc(
                    grad_vars, grad_sharding_specs, mesh_idx,
                    self.schedule.first_backward_batch_index,
                    global_config.use_memzero_for_gradient_accumulation,
                    instruction_lists, executable_config_lists)

        return grad_uuids, instruction_lists

    def _compile_collect_mesh_input(self, mesh_idx, batch_vars):
        mesh_arg_set = OrderedSet()
        var_to_spec = {}
        mesh_batch_vars = OrderedSet()
        num_batch = self.num_batch
        mesh_arg_indices = []
        input_shard_indices = []
        input_shard_specs = []
        mesh_invar_is_batch = []
        for stage_idx in self.schedule.mesh_stage_mapping[mesh_idx]:
            stage = self.stages[stage_idx]
            for spec, invar in zip(stage.input_sharding_specs, stage.invars):
                if invar in self.global_invars:
                    var_to_spec[invar] = spec
                    if invar in batch_vars:
                        # Split batch arg
                        for batch_idx in range(num_batch):
                            mesh_arg_set.add((invar, batch_idx))
                        mesh_batch_vars.add(invar)
                    else:
                        mesh_arg_set.add((invar, 0))
        mesh_arg_list = list(mesh_arg_set)

        for info in mesh_arg_list:
            var, batch_idx = info
            if batch_idx != 0:
                continue

            global_idx = self.global_invars.index(var)
            mesh_arg_indices.append(global_idx)
            mesh_invar_is_batch.append(self.is_batch[global_idx])

            if self.is_batch[global_idx]:
                aval = var.aval
                batch_dim = 0
                new_shape = (num_batch * aval.shape[0],) + aval.shape[1:]
                new_spec = get_microbatch_sharding_spec(var_to_spec[var],
                                                        batch_dim, num_batch)
                input_shard_indices.append(
                    pxla.spec_to_indices(new_shape, new_spec))
                input_shard_specs.append(new_spec)
            else:
                input_shard_indices.append(
                    pxla.spec_to_indices(var.aval.shape, var_to_spec[var]))
                input_shard_specs.append(var_to_spec[var])
        return (mesh_arg_list, mesh_arg_indices, input_shard_indices,
                input_shard_specs, mesh_invar_is_batch)

    def _compile_split_input_to_microbatches(self, global_batch_invar_set):
        """
        Split batch arguments into micro batches.

        The split is like:
        before: a, b, c, d
        after (b, d are batch args and #mb=2): a, b0, b1, c, d0, d1
        """
        donated_invar_set = OrderedSet()
        global_invar_set = OrderedSet(self.global_invars)
        for stage in self.stages:
            for invar, donate in zip(stage.invars, stage.donated_invars):
                if donate and invar in global_invar_set:
                    donated_invar_set.add(invar)
        num_mesh = len(self.mesh_group)
        mesh_arg_lists = [None for _ in range(num_mesh)]

        # Dispatch args to each mesh
        arg_last_use = {}
        donate_invars = []
        mesh_arg_indices = []
        input_shard_indices = []
        input_shard_specs = []
        batch_invars = []
        for mesh_idx in range(num_mesh):
            (mesh_arg_list, arg_indices, shard_indices, shard_specs,
             is_batch) = self._compile_collect_mesh_input(
                 mesh_idx, global_batch_invar_set)

            mesh_arg_lists[mesh_idx] = mesh_arg_list
            delete_after_run = [
                var in donated_invar_set or var in global_batch_invar_set
                for var, _ in mesh_arg_list
            ]
            donate_invars.append(delete_after_run)
            for info in mesh_arg_list:
                var, batch_idx = info
                if batch_idx != 0:
                    continue
                arg_last_use[var] = mesh_idx

            mesh_arg_indices.append(arg_indices)
            input_shard_indices.append(shard_indices)
            input_shard_specs.append(shard_specs)
            batch_invars.append(is_batch)

        delete_after_shard = []
        for mesh_idx in range(num_mesh):
            delete_after_shard.append([
                self.global_invars[idx] in donated_invar_set and
                arg_last_use[self.global_invars[idx]] == mesh_idx
                for idx in mesh_arg_indices[mesh_idx]
            ])

        # Get local uuids for each input
        input_local_uuid_lists = defaultdict(list)
        for mesh_idx, physical_mesh in enumerate(self.mesh_group):
            mesh_arg_list = mesh_arg_lists[mesh_idx]
            num_args = len(mesh_arg_list)
            # shape: (num_args, num_hosts, num_devices_per_host)
            arg_uuids = self._get_next_uuids(
                num_args * physical_mesh.num_devices).reshape(
                    num_args, -1, physical_mesh.num_devices_per_host)
            for arg_idx, info in enumerate(mesh_arg_lists[mesh_idx]):
                var, batch_idx = info
                self.env.set_var_mesh_uuids(var, batch_idx, mesh_idx,
                                            arg_uuids[arg_idx])
                for worker_idx, worker in enumerate(physical_mesh.workers):
                    input_local_uuid_lists[worker].append(arg_uuids[arg_idx,
                                                                    worker_idx])
        input_config = PipeshardInputConfig(
            input_local_uuid_lists=input_local_uuid_lists,
            donate_invars=donate_invars,
            mesh_arg_indices=mesh_arg_indices,
            input_shard_indices=input_shard_indices,
            delete_after_shard=delete_after_shard,
            batch_invars=batch_invars)
        return input_config, input_shard_specs

    def _compile_concate_get_spec(self, to_concate_vars):
        var_to_spec_all_meshes = []
        output_at = defaultdict(OrderedSet)
        num_mesh = len(self.mesh_group)
        for mesh_idx in range(num_mesh):
            var_to_spec = {}
            for stage_idx in self.schedule.mesh_stage_mapping[mesh_idx]:
                stage = self.stages[stage_idx]
                for spec, outvar in zip(stage.output_sharding_specs,
                                        stage.outvars):
                    if outvar in to_concate_vars:
                        var_to_spec[outvar] = spec
                        output_at[outvar].add(mesh_idx)
            var_to_spec_all_meshes.append(var_to_spec)
        return var_to_spec_all_meshes, output_at

    def _compile_concate(self, instruction_lists, executable_config_lists):
        """
        Generate concate instruction for variables used in non-microbatch part,
        but are not reduced. They should be concated.
        """
        batch_dim = 0
        to_concate_vars = set(self.concat_vars_mapping.values())
        to_concate_specs, output_at = self._compile_concate_get_spec(
            to_concate_vars)
        for var in self.concat_vars_mapping:
            src_var = self.concat_vars_mapping[var]
            dst_mesh_to_uuids = self.env.get_var_meshes(
                var, self.schedule.last_backward_batch_index)
            for mesh_idx in output_at[src_var]:
                physical_mesh = self.mesh_group[mesh_idx]
                # Get input and output uuids
                input_args = np.zeros((self.num_batch, *physical_mesh.shape),
                                      dtype=np.int64)
                for batch_idx in range(self.num_batch):
                    input_args[batch_idx] = self.env.get_var_mesh_uuids(
                        src_var, batch_idx, mesh_idx)
                outputs = self._get_next_uuids(
                    physical_mesh.num_devices).reshape(
                        (1, *physical_mesh.shape))
                dst_mesh_to_uuids[mesh_idx] = outputs.reshape(
                    physical_mesh.shape)
                input_args = input_args.transpose([1, 0, 2])
                outputs = outputs.transpose([1, 0, 2])

                # create and run concat executable
                exec_uuid = next_mesh_executable_uuid()
                spec = to_concate_specs[mesh_idx][src_var]
                hlo_proto = compile_concatenate(xb.get_backend("gpu"),
                                                physical_mesh.shape, spec,
                                                self.num_batch, batch_dim,
                                                src_var.aval)
                exec_config = ConcatWorkerExecutableConfig(exec_uuid, hlo_proto)
                kwargs = {
                    "sync_before": False,
                    "sync_after": False,
                }
                for worker_idx, worker in enumerate(physical_mesh.workers):
                    executable_config_lists[worker].append(exec_config)
                    instruction_lists[worker].append(
                        PipelineInstruction.Run(exec_uuid,
                                                input_args[worker_idx],
                                                outputs[worker_idx], kwargs))

    def _compile_collect_outputs(self) -> None:
        """
        Generate output information.

        This function dispatches output information, including local uuid, local indices to global
        indices, and output specs to each mesh.
        """
        num_mesh = len(self.mesh_group)
        output_local_uuid_list = defaultdict(list)  # Dict[worker -> List[uuid]]

        # collect outvar specs
        var_to_spec_all_meshes = []
        global_outvar_set = OrderedSet(self.global_outvars)
        # This is only a patch. It will be deprecated after we move concat into a stage
        reversed_concat = {
            v: k
            for k, v in self.concat_vars_mapping.items()
            if k in global_outvar_set
        }
        output_at = defaultdict(OrderedSet)
        for mesh_idx in range(num_mesh):
            var_to_spec = {}
            for stage_idx in self.schedule.mesh_stage_mapping[mesh_idx]:
                stage = self.stages[stage_idx]
                for spec, outvar in zip(stage.output_sharding_specs,
                                        stage.outvars):
                    if outvar in global_outvar_set:
                        var_to_spec[outvar] = spec
                        output_at[outvar].add(mesh_idx)
                    if outvar in reversed_concat:
                        concat_outvar = reversed_concat[outvar]
                        var_to_spec[concat_outvar] = spec
                        output_at[concat_outvar].add(mesh_idx)
            var_to_spec_all_meshes.append(var_to_spec)
        # assign indices and get specs
        for outvar in self.global_outvars:
            # the apply gradient only writes to microbatch 0
            mesh_to_uuid = self.env.get_var_meshes(
                outvar, self.schedule.last_backward_batch_index)
            mesh_out_indices = {}
            for mesh_idx in output_at[outvar]:
                mesh = self.mesh_group[mesh_idx]
                uuids = mesh_to_uuid[mesh_idx]
                for worker_idx, worker in enumerate(mesh.workers):
                    output_local_uuid_list[worker].append(uuids[worker_idx])
                mesh_out_indices[mesh_idx] = (
                    len(output_local_uuid_list[mesh.workers[0]]) - 1)
                self.output_spec_list[mesh_idx].append(
                    var_to_spec_all_meshes[mesh_idx][outvar])
            self.mesh_output_indices.append(mesh_out_indices)

        return output_local_uuid_list

    def _compile_alloc(self, variables, sharding_specs, mesh_idx, batch_idx,
                       preallocated, instruction_lists,
                       executable_config_lists):
        """Compile an executable which allocates zero buffers.

        The zero buffers are:
        1) gradient accumulation buffers
        2) temp buffers for receiving tensors
        """
        config_class = (MemZeroWorkerExecutableConfig
                        if preallocated else AllocateZeroWorkerExecutableConfig)
        avals = [var.aval for var in variables]
        sharded_shapes = [
            get_shard_shape(aval, spec)
            for aval, spec in zip(avals, sharding_specs)
        ]
        dtypes = [aval.dtype for aval in avals]
        exec_uuid = next_mesh_executable_uuid()
        config = config_class(exec_uuid, sharded_shapes, dtypes)

        physical_mesh = self.mesh_group[mesh_idx]
        output_uuids = self._get_next_uuids(
            len(variables) * physical_mesh.num_devices).reshape(
                len(physical_mesh.workers), len(variables), -1)
        for worker_idx, worker in enumerate(physical_mesh.workers):
            executable_config_lists[worker].append(config)
            if preallocated:
                in_uuids = output_uuids[worker_idx]
                out_uuids = []
            else:
                in_uuids = []
                out_uuids = output_uuids[worker_idx]
            instruction_lists[worker].append(
                PipelineInstruction.Run(config.exec_uuid,
                                        in_uuids,
                                        out_uuids, {
                                            "sync_before": False,
                                            "sync_after": False
                                        },
                                        info="mem set zero" if preallocated else
                                        "allocate zero for recv"))

        # shape: (#args, num_hosts, num_devices_per_host)
        transposed = output_uuids.transpose([1, 0, 2])
        for var_idx, var in enumerate(variables):
            self.env.set_var_mesh_uuids(var, batch_idx, mesh_idx,
                                        transposed[var_idx])
        return output_uuids

    def _get_outs_handler(self):
        """Setup outs handlers that assemble RemoteBufs into DistributedArrays."""
        avals = [outvar.aval for outvar in self.global_outvars]
        is_replicated = [
            bool(len(self.outvar_index_to_mesh_index_mapping[i]) > 1)
            for i, _ in enumerate(self.global_outvars)
        ]

        def outs_handler(mesh_group, bufs):
            ret = []
            for i, aval in enumerate(avals):
                if not is_replicated[i]:
                    # construct DistributedArray
                    mesh_idx = self.outvar_index_to_mesh_index_mapping[i][0]
                    device_mesh = mesh_group[mesh_idx]
                    outvar_index_on_mesh = self.mesh_output_indices[i][mesh_idx]
                    spec = self.output_spec_list[mesh_idx][outvar_index_on_mesh]
                    arr = DistributedArray(
                        device_mesh=device_mesh,
                        aval=aval,
                        sharding_spec=spec,
                        remote_buffers=bufs[mesh_idx][outvar_index_on_mesh],
                        indices=pxla.spec_to_indices(aval.shape, spec))
                else:
                    # otherwise, construct RepliatedDistributedArray
                    meshes = []
                    distributed_arrays = []
                    for _, mesh_idx in enumerate(
                            self.outvar_index_to_mesh_index_mapping[i]):
                        meshes.append(mesh_group[mesh_idx])
                        outvar_index_on_mesh = self.mesh_output_indices[i][
                            mesh_idx]
                        spec = self.output_spec_list[mesh_idx][
                            outvar_index_on_mesh]
                        distributed_arrays.append(
                            DistributedArray(device_mesh=mesh_group[mesh_idx],
                                             aval=aval,
                                             sharding_spec=spec,
                                             remote_buffers=bufs[mesh_idx]
                                             [outvar_index_on_mesh],
                                             indices=pxla.spec_to_indices(
                                                 aval.shape, spec)))
                    arr = ReplicatedDistributedArray(meshes, distributed_arrays)
                ret.append(arr)
            return ret

        return outs_handler

    def _compile_load_info(self, mesh_arg_indices, input_shard_specs):
        assert self.in_tree is not None

        # build load_info_arr: flatten global index => LoadInfo object
        load_info_arr = [None] * len(self.is_batch)
        for mesh_idx, physical_mesh in enumerate(self.mesh_group):
            for local_idx, global_idx in enumerate(mesh_arg_indices[mesh_idx]):
                aval, mesh, spec = (self.global_invars[global_idx].aval,
                                    physical_mesh,
                                    input_shard_specs[mesh_idx][local_idx])
                if load_info_arr[global_idx] is None:
                    load_info_arr[global_idx] = LoadInfo([aval], [mesh], [spec])
                else:
                    load_info_arr[global_idx].add_replica(aval, mesh, spec)

        # build load_info_arr
        return tree_unflatten(self.in_tree, load_info_arr)

    # TODO(yonghao): set empty buffer is not compatiable with local allgather
    @staticmethod
    def _compile_resharding_task(src_mesh,
                                 dst_mesh,
                                 src_uuids: np.ndarray,
                                 resharding_task: SymbolicReshardingTask,
                                 recv_uuids: np.ndarray,
                                 instruction_lists,
                                 set_empty_buffer=False):
        """
        Compile and generate SEND and RECV PipelineInstructions for a ReshardingTask.

        Args:
            src_mesh: the src mesh
            dst_mesh: the dst mesh
            src_uuids: uuids of resharded buffer in src mesh
            resharding_task: the task to be compiled
            recv_uuids: uuids of resharded buffer in dst mesh
            set_empty_buffer: set the empty buffer when recv or not
        """
        num_devices_per_host = dst_mesh.num_devices_per_host
        send_buf_uuids = {worker: [] for worker in src_mesh.workers}
        recv_buf_uuids = {worker: [] for worker in dst_mesh.workers}
        num_device_sender_host = src_mesh.num_devices_per_host

        # collect uuids of each send_tile in each worker based on resharding_task's plan
        for sender_str in resharding_task.sender_uuid_plan:
            send_worker = resharding_task.collective_group.device_str_to_mesh_worker_map[
                sender_str]
            send_buf_flat_idx = resharding_task.task_spec.src.device_str_to_flat_index[
                sender_str]
            send_buf_host = send_buf_flat_idx // num_device_sender_host
            send_buf_device = send_buf_flat_idx % num_device_sender_host
            send_buf_uuids[send_worker].append(src_uuids[send_buf_host,
                                                         send_buf_device])

        # add send tasks for each worker
        for w, task_uuid in resharding_task.send_worker_task_ids.items():
            input_uuids = send_buf_uuids[w]
            instruction_lists[w].append(
                PipelineInstruction.Send(task_uuid, input_uuids))

        # collect uuids of each recv_tile in each worker based on resharding_task's plan
        for receiver_str in resharding_task.receiver_uuid_plan:
            receiver_worker = resharding_task.collective_group.device_str_to_mesh_worker_map[
                receiver_str]
            recv_buf_flat_idx = resharding_task.task_spec.dst.device_str_to_flat_index[
                receiver_str]
            recv_buf_host = recv_buf_flat_idx // num_devices_per_host
            recv_buf_device = recv_buf_flat_idx % num_devices_per_host
            recv_buf_uuids[receiver_worker].append(recv_uuids[recv_buf_host,
                                                              recv_buf_device])

        # add recv task for each worker
        for w, task_uuid in resharding_task.recv_worker_task_ids.items():
            output_uuids = recv_buf_uuids[w]
            allgather_uuid = (resharding_task.allgather_worker_task_ids[w] if
                              resharding_task.is_local_allgather_task else None)
            instruction_lists[w].append(
                PipelineInstruction.Recv(task_uuid, output_uuids,
                                         set_empty_buffer, allgather_uuid))

    @staticmethod
    def _compile_broadcast_resharding_task(
            src_mesh,
            dst_mesh,
            src_uuids: np.ndarray,
            resharding_task: SymbolicBroadcastReshardingTask,
            recv_uuids: np.ndarray,
            instruction_lists,
            set_empty_buffer=False):

        # add broadcast-based resharding task for each worker
        for w, task_uuid in resharding_task.broadcast_worker_task_ids.items():
            output_uuids = None
            input_uuids = None
            if w in src_mesh.workers:
                host_id = src_mesh.workers.index(w)
                input_uuids = src_uuids[host_id]
            else:
                host_id = dst_mesh.workers.index(w)
                output_uuids = recv_uuids[host_id]
            instruction_lists[w].append(
                PipelineInstruction.Broadcast(task_uuid, input_uuids,
                                              output_uuids, "broadcast"))

    @staticmethod
    def _compile_free(worker, used_outside, donated, instruction_lists):
        """Compile and generate FREE PipelineInstruction to recycle memory."""
        instruction_list = instruction_lists[worker]
        new_list = []
        cannot_free_uuids = OrderedSet(used_outside)
        cannot_free_uuids.update(donated)
        for instruction in reversed(instruction_list):
            # for free instruction, do not free again
            if instruction.input_uuids is None:
                new_list.append(instruction)
                continue
            input_uuids = flatten_uuid_set(instruction.input_uuids)
            if not instruction.opcode == PipelineInstType.FREE:
                unused_uuids = input_uuids.difference(cannot_free_uuids)
                if len(unused_uuids) > 0:
                    new_list.append(
                        PipelineInstruction.Free(np.array(list(unused_uuids))))
            cannot_free_uuids.update(input_uuids)
            new_list.append(instruction)
        return list(reversed(new_list))

    @cached_property
    def outvar_index_to_mesh_index_mapping(self) -> Dict[int, List[int]]:
        """
        A mapping from an outvar to the indices of meshes it locates on.

        Returns:
            mapping (Dict[int, List[int]]): mapping[outvar_idx] is a list
                containing the indices of meshes it locates on.
        """
        if self.mesh_output_indices is None:
            raise RuntimeError()
        mapping = {}
        for i, _ in enumerate(self.global_outvars):
            mapping[i] = list(self.mesh_output_indices[i].keys())
        return mapping

max_stars_repo_name: alpa-projects/alpa

max_issues_repo_name: alpa-projects/alpa

max_forks_repo_name: alpa-projects/alpa

reason: import numpy,from jax


text:
import warnings
import rasterio
from rasterio.shutil import copy as rio_copy
import numpy as np
import xarray as xr
import dask
from dask.delayed import Delayed
from pathlib import Path
from typing import Union, Optional, List, Any

from .io import check_write_path
from .geometry import GeoBox
from .geometry.tools import align_up

__all__ = ("write_cog", "to_cog")

def _adjust_blocksize(block, dim):
    if block > dim:
        return align_up(dim, 16)
    return align_up(block, 16)

def _write_cog(pix: np.ndarray,
               geobox: GeoBox,
               fname: Union[Path, str],
               nodata: Optional[float] = None,
               overwrite: bool = False,
               blocksize: Optional[int] = None,
               overview_resampling: Optional[str] = None,
               overview_levels: Optional[List[int]] = None,
               ovr_blocksize: Optional[int] = None,
               **extra_rio_opts) -> Union[Path, bytes]:
    """Write geo-registered ndarray to GeoTiff file or RAM.

    :param pix: ``xarray.DataArray`` with crs or (ndarray, geobox, nodata) triple
    :param fname:  Output file or ":mem:"
    :param nodata: Set ``nodata`` flag to this value if supplied
    :param overwrite: True -- replace existing file, False -- abort with IOError exception
    :param blocksize: Size of internal tiff tiles (512x512 pixels)
    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
    :param overview_resampling: Use this resampling when computing overviews
    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]
                            to disable overviews supply empty list ``[]``
    :param extra_rio_opts: Any other option is passed to ``rasterio.open``

    When fname=":mem:" write COG to memory rather than to a file and return it
    as memoryview object.

    NOTE: about memory requirements

    This function generates temporary in memory tiff file without compression
    to speed things up. It then adds overviews to this file and only then
    copies it to the final destination with requested compression settings.
    This is necessary to produce compliant COG, since COG standard demands
    overviews to be placed before native resolution data and double pass is the
    only way to achieve this currently.

    This means that this function will use about 1.5 to 2 times memory taken by `pix`.
    """
    # pylint: disable=too-many-locals
    if blocksize is None:
        blocksize = 512
    if ovr_blocksize is None:
        ovr_blocksize = blocksize
    if overview_levels is None:
        overview_levels = [2 ** i for i in range(1, 6)]
    if overview_resampling is None:
        overview_resampling = "nearest"

    if pix.ndim == 2:
        h, w = pix.shape
        nbands = 1
        band = 1  # type: Any
    elif pix.ndim == 3:
        if pix.shape[:2] == geobox.shape:
            pix = pix.transpose([2, 0, 1])
        elif pix.shape[-2:] != geobox.shape:
            raise ValueError('GeoBox shape does not match image shape')

        nbands, h, w = pix.shape
        band = tuple(i for i in range(1, nbands + 1))
    else:
        raise ValueError("Need 2d or 3d ndarray on input")

    assert geobox.shape == (h, w)

    if fname != ":mem:":
        path = check_write_path(
            fname, overwrite
        )  # aborts if overwrite=False and file exists already

    resampling = rasterio.enums.Resampling[overview_resampling]

    if (blocksize % 16) != 0:
        warnings.warn(f"Block size must be a multiple of 16, will be adjusted")

    rio_opts = dict(
        width=w,
        height=h,
        count=nbands,
        dtype=pix.dtype.name,
        crs=str(geobox.crs),
        transform=geobox.transform,
        tiled=True,
        blockxsize=_adjust_blocksize(blocksize, w),
        blockysize=_adjust_blocksize(blocksize, h),
        zlevel=6,
        predictor=3 if pix.dtype.kind == "f" else 2,
        compress="DEFLATE",
    )

    if nodata is not None:
        rio_opts.update(nodata=nodata)

    rio_opts.update(extra_rio_opts)

    # Deal efficiently with "no overviews needed case"
    if len(overview_levels) == 0:
        if fname == ":mem:":
            with rasterio.MemoryFile() as mem:
                with mem.open(driver="GTiff", **rio_opts) as dst:
                    dst.write(pix, band)
                return bytes(mem.getbuffer())
        else:
            with rasterio.open(path, mode='w', driver='GTiff', **rio_opts) as dst:
                dst.write(pix, band)
            return path

    # copy re-compresses anyway so skip compression for temp image
    tmp_opts = rio_opts.copy()
    tmp_opts.pop("compress")
    tmp_opts.pop("predictor")
    tmp_opts.pop("zlevel")

    with rasterio.Env(GDAL_TIFF_OVR_BLOCKSIZE=ovr_blocksize):
        with rasterio.MemoryFile() as mem:
            with mem.open(driver="GTiff", **tmp_opts) as tmp:
                tmp.write(pix, band)
                tmp.build_overviews(overview_levels, resampling)

                if fname == ":mem:":
                    with rasterio.MemoryFile() as mem2:
                        rio_copy(
                            tmp,
                            mem2.name,
                            driver="GTiff",
                            copy_src_overviews=True,
                            **rio_opts
                        )
                        return bytes(mem2.getbuffer())

                rio_copy(
                    tmp, path, driver="GTiff", copy_src_overviews=True, **rio_opts
                )

    return path

_delayed_write_cog_to_mem = dask.delayed(  # pylint: disable=invalid-name
    _write_cog,
    name="compress-cog", pure=True, nout=1
)

_delayed_write_cog_to_file = dask.delayed(   # pylint: disable=invalid-name
    _write_cog,
    name="save-cog", pure=False, nout=1
)

def write_cog(geo_im: xr.DataArray,
              fname: Union[str, Path],
              blocksize: Optional[int] = None,
              ovr_blocksize: Optional[int] = None,
              overview_resampling: Optional[str] = None,
              overview_levels: Optional[List[int]] = None,
              **extra_rio_opts) -> Union[Path, bytes, Delayed]:
    """
    Save ``xarray.DataArray`` to a file in Cloud Optimized GeoTiff format.

    This function is "Dask aware". If ``geo_im`` is a Dask array, then the
    output of this function is also a Dask Delayed object. This allows us to
    save multiple images concurrently across a Dask cluster. If you are not
    familiar with Dask this can be confusing, as no operation is performed until
    ``.compute()`` method is called, so if you call this function with Dask
    array it will return immediately without writing anything to disk.

    If you are using Dask to speed up data loading, follow example below:

    .. code-block:: python

       # Example: save red band from first time slice to file "red.tif"
       xx = dc.load(.., dask_chunks=dict(x=1024, y=1024))
       write_cog(xx.isel(time=0).red, "red.tif").compute()
       # or compute input first instead
       write_cog(xx.isel(time=0).red.compute(), "red.tif")

    :param geo_im: ``xarray.DataArray`` with crs
    :param fname: Output path or ``":mem:"`` in which case compress to RAM and return bytes
    :param blocksize: Size of internal tiff tiles (512x512 pixels)
    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
    :param overview_resampling: Use this resampling when computing overviews
    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32],
                            to disable overviews supply empty list ``[]``
    :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is
                   read from the attributes of the input array (``geo_im.attrs['nodata']``).
    :param extra_rio_opts: Any other option is passed to ``rasterio.open``

    :returns: Path to which output was written
    :returns: Bytes if ``fname=":mem:"``
    :returns: ``dask.Delayed`` object if input is a Dask array

    .. note ::

       **memory requirements**

       This function generates temporary in memory tiff file without
       compression to speed things up. It then adds overviews to this file and
       only then copies it to the final destination with requested compression
       settings. This is necessary to produce compliant COG, since COG standard
       demands overviews to be placed before native resolution data and double
       pass is the only way to achieve this currently.

       This means that this function will use about 1.5 to 2 times memory taken by ``geo_im``.
    """
    pix = geo_im.data
    geobox = getattr(geo_im, 'geobox', None)
    nodata = extra_rio_opts.pop('nodata', None)
    if nodata is None:
        nodata = geo_im.attrs.get("nodata", None)

    if geobox is None:
        raise ValueError("Need geo-registered array on input")

    if dask.is_dask_collection(pix):
        real_op = _delayed_write_cog_to_mem if fname == ":mem:" else _delayed_write_cog_to_file
    else:
        real_op = _write_cog

    return real_op(
        pix,
        geobox,
        fname,
        nodata=nodata,
        blocksize=blocksize,
        ovr_blocksize=ovr_blocksize,
        overview_resampling=overview_resampling,
        overview_levels=overview_levels,
        **extra_rio_opts)

def to_cog(geo_im: xr.DataArray,
           blocksize: Optional[int] = None,
           ovr_blocksize: Optional[int] = None,
           overview_resampling: Optional[str] = None,
           overview_levels: Optional[List[int]] = None,
           **extra_rio_opts) -> Union[bytes, Delayed]:
    """
    Compress ``xarray.DataArray`` into Cloud Optimized GeoTiff bytes in memory.

    This function doesn't write to disk, it compresses in RAM, this is useful
    for saving data to S3 or other cloud object stores.

    This function is "Dask aware". If ``geo_im`` is a Dask array, then the
    output of this function is also a Dask Delayed object. This allows us to
    compress multiple images concurrently across a Dask cluster. If you are not
    familiar with Dask this can be confusing, as no operation is performed until
    ``.compute()`` method is called, so if you call this function with Dask
    array it will return immediately without compressing any data.

    :param geo_im: ``xarray.DataArray`` with crs
    :param blocksize: Size of internal tiff tiles (512x512 pixels)
    :param ovr_blocksize: Size of internal tiles in overview images (defaults to blocksize)
    :param overview_resampling: Use this resampling when computing overviews
    :param overview_levels: List of shrink factors to compute overiews for: [2,4,8,16,32]
    :param nodata: Set ``nodata`` flag to this value if supplied, by default ``nodata`` is
                   read from the attributes of the input array (``geo_im.attrs['nodata']``).
    :param extra_rio_opts: Any other option is passed to ``rasterio.open``

    :returns: In-memory GeoTiff file as bytes
    :returns: ``dask.Delayed`` object if input is a Dask array

    Also see :py:meth:`~datacube.utils.cog.write_cog`

    """
    bb = write_cog(geo_im,
                   ":mem:",
                   blocksize=blocksize,
                   ovr_blocksize=ovr_blocksize,
                   overview_resampling=overview_resampling,
                   overview_levels=overview_levels,
                   **extra_rio_opts)

    assert isinstance(bb, (bytes, Delayed))  # for mypy sake for :mem: output it bytes or delayed bytes
    return bb

ext: py

reason: import numpy


text:
import logging
import pathlib
from typing import Union, Optional, Tuple, List

import numpy as np

from pylas.vlrs.known import ExtraBytesStruct, ExtraBytesVlr
from .. import extradims, errors
from ..compression import LazBackend
from ..laswriter import LasWriter
from ..point import record, dims, PointFormat
from ..point.dims import ScaledArrayView
from ..vlrs import vlrlist

logger = logging.getLogger(__name__)

class LasBase(object):
    """LasBase is the base of all the different LasData classes.
    These classes are objects that the user will interact with to manipulate las data.

    It connects the point record, header, vlrs together.

    To access points dimensions using this class you have two possibilities

    .. code:: python

        las = pylas.read('some_file.las')
        las.classification
        # or
        las['classification']

    .. note::
        using las['dimension_name']  is not possible with the scaled values of x, y, z

    """

    def __init__(self, *, header, vlrs=None, points=None):
        if points is None:
            points = record.PackedPointRecord.empty(PointFormat(header.point_format_id))
        self.__dict__["_points"] = points
        self.header = header
        self.vlrs = vlrs if vlrs is not None else vlrlist.VLRList()

    @property
    def x(self):
        """Returns the scaled x positions of the points as doubles"""
        return ScaledArrayView(self.X, self.header.x_scale, self.header.x_offset)

    @x.setter
    def x(self, value):
        if len(value) > len(self.points):
            self.points.resize(len(value))
        self.x[:] = value

    @property
    def y(self):
        """Returns the scaled y positions of the points as doubles"""
        return ScaledArrayView(self.Y, self.header.y_scale, self.header.y_offset)

    @y.setter
    def y(self, value):
        if len(value) > len(self.points):
            self.points.resize(len(value))
        self.y[:] = value

    @property
    def z(self):
        """Returns the scaled z positions of the points as doubles"""
        return ScaledArrayView(self.Z, self.header.z_scale, self.header.z_offset)

    @z.setter
    def z(self, value):
        if len(value) > len(self.points):
            self.points.resize(len(value))
        self.z[:] = value

    @property
    def point_format(self):
        return self.points.point_format

    @property
    def points(self):
        return self._points

    @points.setter
    def points(self, new_points):
        if new_points.point_format != self._points.point_format:
            raise errors.IncompatibleDataFormat(
                "Cannot set points with a different point format, convert first"
            )
        self._points = new_points
        self.update_header()

    def change_scaling(self, scales=None, offsets=None) -> None:
        if scales is None:
            scales = self.header.scales
        if offsets is None:
            offsets = self.header.offsets

        record.apply_new_scaling(self, scales, offsets)

        self.header.scales = scales
        self.header.offsets = offsets

    def __getattr__(self, item):
        """Automatically called by Python when the attribute
        named 'item' is no found. We use this function to forward the call the
        point record. This is the mechanism used to allow the users to access
        the points dimensions directly through a LasData.

        Parameters
        ----------
        item: str
            name of the attribute, should be a dimension name

        Returns
        -------
        The requested dimension if it exists

        """
        try:
            return self.points[item]
        except ValueError:
            raise AttributeError(
                f"{self.__class__.__name__} object has no attribute '{item}'"
            ) from None

    def __setattr__(self, key, value):
        """This is called on every access to an attribute of the instance.
        Again we use this to forward the call the the points record

        But this time checking if the key is actually a dimension name
        so that an error is raised if the user tries to set a valid
        LAS dimension even if it is not present in the field.
        eg: user tries to set the red field of a file with point format 0:
        an error is raised

        """
        if key in self.point_format.dimension_names:
            self.points[key] = value
        elif key in dims.DIMENSIONS_TO_TYPE:
            raise ValueError(
                f"Point format {self.point_format} does not support {key} dimension"
            )
        else:
            super().__setattr__(key, value)

    def __getitem__(self, item):
        return self.points[item]

    def __setitem__(self, key, value):
        self.points[key] = value

    def add_extra_dim(self, name: str, type: str, description: str = ""):
        """Adds a new extra dimension to the point record

        .. note::

            If you plan on adding multiple extra dimensions,
            prefer :meth:`pylas.LasBase.add_extra_dims` as it will
            save re-allocations and data copy

        Parameters
        ----------
        name: str
            the name of the dimension, spaces are replaced with '_'.
        type: str
            type of the dimension (eg 'uint8' or 'u1')
        description: str, optional
            a small description of the dimension
        """
        self.add_extra_dims([(name, type, description)])

    def add_extra_dims(self, type_tuples: List[Tuple[str,...]]):
        """ Add multiple extra dimensions at once

        Parameters
        ----------

        type_tuples:
               a list of tuple describing the dimensions to add
               [(name, type, description), (name2, other_type)]
               The description is optional
        """
        extra_bytes_structs = []
        tuples = []
        for name, type, *rest in type_tuples:
            name = name.replace(" ", "_")
            if rest:
                description = rest[0]
            else:
                description = ""
            type_id = extradims.get_id_for_extra_dim_type(type)
            tuples.append((name, extradims.get_type_for_extra_dim(type_id)))
            extra_bytes_structs.append(ExtraBytesStruct(
                data_type=type_id, name=name.encode(), description=description.encode()
            ))

        try:
            extra_bytes_vlr = self.vlrs.get("ExtraBytesVlr")[0]
        except IndexError:
            extra_bytes_vlr = ExtraBytesVlr()
            self.vlrs.append(extra_bytes_vlr)
        finally:
            extra_bytes_vlr.extra_bytes_structs.extend(extra_bytes_structs)
            self.points.add_extra_dims(tuples)

    def update_header(self):
        """Update the information stored in the header
        to be in sync with the actual data.

        This method is called automatically when you save a file using
        :meth:`pylas.lasdatas.base.LasBase.write`
        """
        self.header.point_format_id = self.points.point_format.id
        self.header.point_count = len(self.points)
        self.header.point_data_record_length = self.points.point_size

        if len(self.points) > 0:
            self.header.x_max = self.x.max()
            self.header.y_max = self.y.max()
            self.header.z_max = self.z.max()

            self.header.x_min = self.x.min()
            self.header.y_min = self.y.min()
            self.header.z_min = self.z.min()

            unique, counts = np.unique(self.return_number, return_counts=True)
            self.header.number_of_points_by_return = counts

    def write_to(
        self, out_stream, do_compress=False, laz_backend=LazBackend.detect_available()
    ):
        """writes the data to a stream

        Parameters
        ----------
        out_stream: file object
            the destination stream, implementing the write method
        do_compress: bool, optional, default False
            Flag to indicate if you want the date to be compressed
        laz_backend: optional, the laz backend to use
            By default, pylas detect available backends
        """
        with LasWriter(
            out_stream,
            self.header,
            self.vlrs,
            do_compress=do_compress,
            closefd=False,
            laz_backend=laz_backend,
        ) as writer:
            writer.write(self.points)

    @staticmethod
    def _raise_if_not_expected_pos(stream, expected_pos):
        if not stream.tell() == expected_pos:
            raise RuntimeError(
                "Writing, expected to be at pos {} but stream is at pos {}".format(
                    expected_pos, stream.tell()
                )
            )

    def write_to_file(
        self, filename: Union[str, pathlib.Path], do_compress: Optional[bool] = None
    ) -> None:
        """Writes the las data into a file

        Parameters
        ----------
        filename : str
            The file where the data should be written.
        do_compress: bool, optional, default None
            if None the extension of the filename will be used
            to determine if the data should be compressed
            otherwise the do_compress flag indicate if the data should be compressed
        """
        is_ext_laz = pathlib.Path(filename).suffix.lower() == ".laz"
        if is_ext_laz and do_compress is None:
            do_compress = True

        with open(filename, mode="wb+") as out:
            self.write_to(out, do_compress=do_compress)

    def write(
        self, destination, do_compress=None, laz_backend=LazBackend.detect_available()
    ):
        """Writes to a stream or file

        When destination is a string, it will be interpreted as the path were the file should be written to,
        also if do_compress is None, the compression will be guessed from the file extension:

        - .laz -> compressed
        - .las -> uncompressed

        .. note::

            This means that you could do something like:
                # Create .laz but not compressed

                las.write('out.laz', do_compress=False)

                # Create .las but compressed

                las.write('out.las', do_compress=True)

            While it should not confuse Las/Laz readers, it will confuse humans so avoid doing it

        Parameters
        ----------
        destination: str or file object
            filename or stream to write to
        do_compress: bool, optional
            Flags to indicate if you want to compress the data
        laz_backend: optional, the laz backend to use
            By default, pylas detect available backends
        """
        if isinstance(destination, (str, pathlib.Path)):
            self.write_to_file(destination)
        else:
            if do_compress is None:
                do_compress = False
            self.write_to(destination, do_compress=do_compress, laz_backend=laz_backend)

    def __repr__(self):
        return "<LasData({}.{}, point fmt: {}, {} points, {} vlrs)>".format(
            self.header.version_major,
            self.header.version_minor,
            self.points.point_format,
            len(self.points),
            len(self.vlrs),
        )


text:
import pickle
import tensorflow as tf
import numpy as np
from termcolor import cprint

import os
import model
import utils
import time

flags = tf.app.flags
flags.DEFINE_integer("nb_epochs", 100000, "Epoch to train [100 000]")
flags.DEFINE_integer("nb_iter_per_epoch", 100, "Epoch to train [100]")

flags.DEFINE_integer("out_frequency", 200, "Output frequency")
flags.DEFINE_float("learning_rate", 0.0002, "Learning rate of for adam [0.0001")
flags.DEFINE_float("decay_learning_rate_step", 10000, "Step to decay the learning rate [10000]")
flags.DEFINE_float("learning_rate_decay_factor", 0.96, "Learning rate decay [0.96]")
flags.DEFINE_float("max_gradient_norm", 5.0, "Clip gradients to this norm.")

flags.DEFINE_float("beta1", 0.5, "Momentum term of adam [0.5]")
flags.DEFINE_integer("batch_size", 64, "The size of batch images [64]")
flags.DEFINE_integer("vocab_size", 55, "The size of the vocabulary [64]")
flags.DEFINE_float("keep_prob", 0.9, "Dropout ratio [0.5]")

flags.DEFINE_integer("hidden_size", 256, "Hidden size of RNN cell [128]")
flags.DEFINE_integer("num_layers", 1, "Num of layers [1]")

FLAGS = flags.FLAGS
model.FLAGS = FLAGS

s1 = time.time()
# Load the idx_to_chars dictionary
with open(os.path.join('..', 'Data', 'MovieQA', 'idx_to_chars.pkl'), 'rb') as f:
    idx_to_chars = pickle.load(f)

# Load the chars_to_idx dictionary
with open(os.path.join('..', 'Data', 'MovieQA', 'chars_to_idx.pkl'), 'rb') as f:
    chars_to_idx = pickle.load(f)

# Instanciates and build the model for feedforward only
seq2seq = model.Seq2Seq(buckets=[(25, 25)], forward_only=True)
seq2seq.build()

# Restore the trained model's parameters from checkpoint file
sess = tf.Session()
saver, summary_writer = utils.restore(seq2seq, sess, save_name="model_saved/")

# ----------------------------------------------------------------------------------------------
print("PREDICTING")
# User asks something...
question_string = "Hello Jack!"
answer_string = "Hey Steve!"
q = utils.encrypt_single(question_string, chars_to_idx)
a = utils.encrypt_single(answer_string, chars_to_idx)

# Equivalent to utils.get_batch but for one example
#   Prepare the batch (batch_size = 1)
buckets = [(25, 25)]
bucket_id = 0
q_pads = np.zeros([1, buckets[bucket_id][0]])
a_pads = np.zeros([1, buckets[bucket_id][1]])

q_pads[0][:q.shape[0]] = q
a_pads[0][:a.shape[0]] = a

# Processing
outputs, questions, answers = seq2seq.predict(bucket_id=0, session=sess, questions=q_pads, answers=a_pads)

# The model replies
outputs = np.squeeze(outputs)
outputs = np.argmax(outputs, axis=1)
output_string = utils.decrypt_single(list(outputs), idx_to_chars)

questions = np.squeeze(questions)
q_string = utils.decrypt_single(list(questions), idx_to_chars)

answers = np.squeeze(answers)
a_string = utils.decrypt_single(list(answers), idx_to_chars)

cprint("Q : " + q_string.split("<PAD>")[0], color="green")
cprint("A : " + a_string.split("<PAD>")[0], color="yellow")
cprint("O : " + output_string, color="red")
print(int(time.time() - s1), " seconds")


text:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler as lrsch
from tqdm import tqdm
import numpy as np
from sklearn.mixture import GaussianMixture

from .estimator import TrainerBase
from . import block
from . import loss
from . import layers
from .utils import rsample_normal, rsample_categorical, temperature_func1

def gamma_func(e, max_value, minimum_epoch):
    return max((max_value - 1) * (1 - e / minimum_epoch) + 1, 1)

class VAEBase(TrainerBase):

    def get_optimizers(self):
        gen_optimizer = optim.Adam(self.parameters(), lr=self._lr)
        if self._kwargs.get("use_lrsch", False):
            scheduler = lrsch.CosineAnnealingWarmRestarts(gen_optimizer,
                                                          100, 2, 0.0001)
            return [gen_optimizer, scheduler]
        else:
            return [gen_optimizer]

    def _batch_end(self):
        if self._kwargs.get("use_lrsch", False):
            self._ops[-1].step(self._e + self._i / len(self._data_iter))

    def _epoch_start(self):
        self._alpha = 1.
        self._gamma = gamma_func(self._e,
                                 self._gamma_max_value,
                                 self._gamma_min_epoch)
        self._temp = temperature_func1(self._e, 0.01, 0.3, self._epoch)

class MClusterVAEsExact(VAEBase):

    def __init__(
        self, inpts, omic_oupts, n_clusters, omic_zdims, att_method="none",
        omic_hiddens=(), enc_hiddens=[200], enc_z_hiddens=[200],
        dec_hiddens=[200], gamma_max_value=5.0, gamma_min_epoch=300
    ):
        super().__init__()
        self._n_clusters = n_clusters
        self._inpts = inpts
        self._att_method = att_method
        self._gamma_max_value = gamma_max_value
        self._gamma_min_epoch = gamma_min_epoch

        self.encoder_y = block.Encoder(inpts, omic_oupts, [n_clusters],
                                       omic_hiddens, enc_hiddens, att_method)
        self.encoders_z = nn.ModuleList([
            block.VAEncoder(i+n_clusters, o, enc_z_hiddens)
            for i, o in zip(inpts, omic_zdims)
        ])
        self.decoders = nn.ModuleList([
            block.VADncoder(i+n_clusters, o, dec_hiddens)
            for i, o in zip(omic_zdims, inpts)
        ])
        self.z_maps = nn.ModuleList([
            nn.Sequential(
                nn.Linear(n_clusters, o * 2, bias=False),
                layers.SplitModule(o, o)
            ) for o in omic_zdims
        ])

    def get_criterions(self):
        fit_loss = loss.MOVBCExactLoss(self._kwargs["recon_ws"])
        return fit_loss

    def _batch_train(self, xs):
        # p(z_j|y), prepare
        iden = torch.eye(self._n_clusters,
                         device=self._device,
                         dtype=torch.float32)
        zy_mu, zy_logsigma2 = [], []
        for zmap in self.z_maps:
            zy_mu_j, zy_logsigma2_j = zmap(iden)   # ncluster x omic_dim
            zy_mu.append(zy_mu_j)
            zy_logsigma2.append(zy_logsigma2_j)

        # p(y|x_1,...,x_k), encode
        logits = self.encoder_y(xs)

        # p(z_ij|x_j,y=i), i=1,...,n_clusters, j=1,...,k
        z_mu, z_logsigma2, xs_rec = [], [], []
        for i in range(logits.size(1)):
            yi = torch.zeros_like(logits)
            yi[:, i] = 1.
            z_mu_i, z_logsigma2_i, xs_rec_i = [], [], []
            for x, enc, dec in zip(xs, self.encoders_z, self.decoders):
                x_y = torch.cat([x, yi], dim=1)
                z_mu_ij, z_logsigma2_ij = enc(x_y)
                z_ij = rsample_normal(z_mu_ij, z_logsigma2_ij)
                x_ij = dec(torch.cat([yi, z_ij], dim=1))
                z_mu_i.append(z_mu_ij)
                z_logsigma2_i.append(z_logsigma2_ij)
                xs_rec_i.append(x_ij)
            z_mu.append(z_mu_i)
            z_logsigma2.append(z_logsigma2_i)
            xs_rec.append(xs_rec_i)

        loss, losses = self._criterions(xs, xs_rec, logits,
                                        zy_mu, zy_logsigma2,
                                        z_mu, z_logsigma2,
                                        self._alpha, self._gamma)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder_y.to(device)
        self.encoder_y.eval()
        logits = []
        z_mu = [[] for _ in range(dat.n)]
        z_logsigma2 = [[] for _ in range(dat.n)]
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                logit = self.encoder_y(xs)
                y = torch.eye(logit.size(1)).to(logit)[logit.argmax(dim=1)]
                for i, (x, enc) in enumerate(zip(xs, self.encoders_z)):
                    z_mu_j, z_logsigma2_j = enc(torch.cat([x, y], dim=1))
                    z_mu[i].append(z_mu_j)
                    z_logsigma2[i].append(z_logsigma2_j)
                logits.append(logit)

        logits = torch.cat(logits).detach().cpu().numpy()
        z_mu = [torch.cat(z_mu_j).detach().cpu().numpy() for z_mu_j in z_mu]
        z_lgs = [torch.cat(z_lgs_j).detach().cpu().numpy()
                 for z_lgs_j in z_logsigma2]

        return (logits, z_mu, z_lgs)

    def cluster(self, transformed_res):
        return np.argmax(transformed_res[0], axis=1)+1, transformed_res[0]

class MClusterVAEsGumbel(MClusterVAEsExact):

    def get_criterions(self):
        fit_loss = loss.MOVBCGumbelLoss(self._kwargs["recon_ws"])
        return fit_loss

    def _batch_train(self, xs):
        # p(y|x_1,...,x_k), encode
        logits = self.encoder_y(xs)

        y = rsample_categorical(logits, self._temp)

        # p(z_j|y), prepare
        zy_mu, zy_logsigma2 = [], []
        for zmap in self.z_maps:
            zy_mu_j, zy_logsigma2_j = zmap(y)   # ncluster x omic_dim
            zy_mu.append(zy_mu_j)
            zy_logsigma2.append(zy_logsigma2_j)

        # p(z_ij|x_j,y=i), i=1,...,n_clusters, j=1,...,k
        # z_mus, z_logsigma2s, zy_mus, zy_logsigma2s, xss = [], [], [], [], []
        z_mu, z_logsigma2, xs_rec = [], [], []
        for x, enc, dec in zip(xs, self.encoders_z, self.decoders):
            x_y = torch.cat([x, y], dim=1)
            z_mu_ij, z_logsigma2_ij = enc(x_y)
            z_ij = rsample_normal(z_mu_ij, z_logsigma2_ij)
            x_ij = dec(torch.cat([y, z_ij], dim=1))
            z_mu.append(z_mu_ij)
            z_logsigma2.append(z_logsigma2_ij)
            xs_rec.append(x_ij)

        loss, losses = self._criterions(xs, xs_rec, logits,
                                        zy_mu, zy_logsigma2,
                                        z_mu, z_logsigma2,
                                        self._alpha, self._gamma)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder_y.to(device)
        self.encoder_y.eval()
        logits = []
        z_mu = [[] for _ in range(dat.n)]
        z_logsigma2 = [[] for _ in range(dat.n)]
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                logit = self.encoder_y(xs)
                y = torch.eye(logit.size(1)).to(logit)[logit.argmax(dim=1)]
                for i, (x, enc) in enumerate(zip(xs, self.encoders_z)):
                    z_mu_j, z_logsigma2_j = enc(torch.cat([x, y], dim=1))
                    z_mu[i].append(z_mu_j)
                    z_logsigma2[i].append(z_logsigma2_j)
                logits.append(logit)

        logits = torch.cat(logits).detach().cpu().numpy()
        z_mu = [torch.cat(z_mu_j).detach().cpu().numpy() for z_mu_j in z_mu]
        z_lgs = [torch.cat(z_lgs_j).detach().cpu().numpy()
                 for z_lgs_j in z_logsigma2]

        return (logits, z_mu, z_lgs)

    def cluster(self, transformed_res):
        return np.argmax(transformed_res[0], axis=1)+1, transformed_res[0]

class AE(VAEBase):

    def __init__(
        self, inpts, omic_oupts, n_clusters, con_dim, att_method="none",
        omic_hiddens=(), enc_hiddens=[200], dec_hiddens=[200]
    ):
        super().__init__()
        self._n_clusters = n_clusters
        self._inpts = inpts
        self._att_method = att_method

        self.encoder = block.Encoder(inpts, omic_oupts, [con_dim],
                                     omic_hiddens, enc_hiddens, att_method)
        self.decoder = block.Decoder(con_dim, inpts, dec_hiddens)

    def get_criterions(self):
        fit_loss = loss.MultiMSELoss(self._kwargs["recon_ws"], "sum")
        return fit_loss

    def _batch_train(self, xs):
        z = self.encoder(xs)
        xs_rec = self.decoder(z)

        loss, losses = self._criterions(xs, xs_rec)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder.to(device)
        self.encoder.eval()
        zs = []
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                z = self.encoder(xs)
                zs.append(z)

        zs = torch.cat(zs).detach().cpu().numpy()

        return zs

    def cluster(self, transformed_res):
        gmm = GaussianMixture(n_components=self._n_clusters,
                              covariance_type="diag")
        cluster_res = gmm.fit_predict(transformed_res) + 1
        return cluster_res, None

class VAE(VAEBase):

    def __init__(
        self, inpts, omic_oupts, n_clusters, con_dim, att_method="none",
        omic_hiddens=(), enc_hiddens=[200], dec_hiddens=[200]
    ):
        super().__init__()
        self._n_clusters = n_clusters
        self._inpts = inpts
        self._att_method = att_method

        self.encoder = block.Encoder(inpts, omic_oupts, [con_dim, con_dim],
                                     omic_hiddens, enc_hiddens, att_method)
        self.decoder = block.Decoder(con_dim, inpts, dec_hiddens)

    def get_criterions(self):
        fit_loss = loss.VAELoss(self._kwargs["recon_ws"], "sum")
        return fit_loss

    def _batch_train(self, xs):
        z_mu, z_logsigma2 = self.encoder(xs)
        z = rsample_normal(z_mu, z_logsigma2)
        xs_rec = self.decoder(z)

        loss, losses = self._criterions(xs, xs_rec, z_mu, z_logsigma2)
        self._ops[0].zero_grad()
        loss.backward()
        self._ops[0].step()
        losses["total"] = loss
        self._accum.add(bs=self._bs, **losses)

    def transform(self, dat, bs, nw, device, disable=False):
        dataloader = dat.batch_iterator_loader(bs, nw, False, False, True)
        device = torch.device(device)
        self.encoder.to(device)
        self.encoder.eval()
        z_mus, z_logsigma2s = [], []
        with torch.no_grad():
            for xs in tqdm(dataloader, "Embed: ", disable=disable):
                xs = [x.to(device) for x in xs]
                z_mu, z_logsigma2 = self.encoder(xs)
                z_mus.append(z_mu)
                z_logsigma2s.append(z_logsigma2)

        z_mus = torch.cat(z_mus).detach().cpu().numpy()
        z_logsigma2s = torch.cat(z_logsigma2s).detach().cpu().numpy()

        return z_mus, z_logsigma2s

    def cluster(self, transformed_res):
        gmm = GaussianMixture(n_components=self._n_clusters,
                              covariance_type="diag")
        cluster_res = gmm.fit_predict(transformed_res[0]) + 1
        return cluster_res, None

max_stars_repo_name: luyiyun/MCluster-VAEs

max_issues_repo_name: luyiyun/MCluster-VAEs

max_forks_repo_name: luyiyun/MCluster-VAEs


text:
from scipy.optimize import minimize
import pylab as plt
import numpy as np
import random

# fit two blackbodies to the synthetic data

def funcinv(wa,D,f,T1):
    kv=k850*(850/wa)**2
    Md=f/(kv/D**2*blackbody_lam(wa, T1)*2.08*10**11)#*1.98*10**30/9.5/10**32/10**12*10**26
    return np.log10(Md)

for i in range(len(data[:,0])):

	ydata = [data[i,30],data[i,34],data[i,36],data[i,38],data[i,40],data[i,42]]
	yerr = [data[i,31],data[i,35],data[i,37],data[i,39],data[i,41],data[i,43]]
	D = data[i,1]*3*10**5/67.30

	chi2 = float('inf')
	T1dist = np.arange(10,30,3)
	T2dist = np.arange(30,60,10)
	Mddist = np.arange(4,6,0.02)
	Mdprob = np.zeros(len(Mddist))
	ptot = 0
	for ii in range(len(Mddist)):
		for T2 in T2dist:
			for T1 in T1dist:
				popt = minimize(leastsq, [0.05], args=(Mddist[ii],wa,ydata,yerr,D,T1,T2),method='Nelder-Mead',options={'maxiter':200})
				Mdr_new = popt.x
				chi2_new = leastsq(Mdr_new,Mddist[ii],wa,ydata,yerr,D,T1,T2)
				if chi2 > chi2_new:
					chi2 = chi2_new
					Mdratio_sav = Mdr_new
					T1_sav = T1
					T2_sav = T2
					Md_sav = Mddist[ii]
					print Mddist[ii],Mdr_new,T1_sav,T2_sav
				prob=np.exp(-0.5*chi2_new)
				Mdprob[ii]+=prob
				ptot+=prob
	Mdprob = Mdprob / ptot
	print "tcold", Mddist[np.where(Mdprob==max(Mdprob))], percentiles(Mddist,Mdprob,16),percentiles(Mddist,Mdprob,50),percentiles(Mddist,Mdprob,84)
	plt.plot(Mddist, Mdprob)
	plt.title(dataID[i])
	plt.xlabel('T_cold (K)')
	plt.show()
	print Md_sav, T1_sav, T2_sav
	plt.errorbar(wa, ydata, yerr=yerr, fmt='bo')
	x2=np.arange(10.,600.,0.1)
	plt.loglog()
	plt.title(dataID[i])
	plt.ylim(0.001,1)
	plt.plot(x2, func2(x2,D,np.log10(Mdratio_sav)+Md_sav,T1_sav,np.log10(Mdratio_sav)+Md_sav,T2_sav), 'r', label="best fit")
	plt.xlabel('Wavelength (microns)')
	plt.ylabel('Flux (Jy)')
	plt.plot(x2, func(x2,D,np.log10(Mdratio_sav)+Md_sav,T1_sav), ':', lw=2,label="cold dust: logMd = %s, Tc= %s K "%(np.log10(Mdratio_sav)+Md_sav,T1_sav) )
	plt.plot(x2, func(x2,D,np.log10(Mdratio_sav)+Md_sav,T2_sav), ':', lw=2,label="warm dust: logMd = %s, Tc= %s K "%(np.log10(Mdratio_sav)+Md_sav,T2_sav) )
	plt.legend(loc=4)
	plt.show()

plt.legend(frameon=False)
plt.savefig('fit_bb.png')

max_stars_repo_name: wdobbels/CAAPR

max_issues_repo_name: wdobbels/CAAPR

max_forks_repo_name: wdobbels/CAAPR

reason: import numpy,from scipy


text:
import os
import sys
import cv2
import numpy as np

def get_binary_img(img):
    # gray img to bin image
    bin_img = np.zeros(shape=(img.shape), dtype=np.uint8)
    h = img.shape[0]
    w = img.shape[1]
    for i in range(h):
        for j in range(w):
            bin_img[i][j] = 255 if img[i][j] > 127 else 0
    return bin_img

def get_contour(bin_img):
    # get contour
    contour_img = np.zeros(shape=(bin_img.shape),dtype=np.uint8)
    contour_img += 255
    h = bin_img.shape[0]
    w = bin_img.shape[1]
    for i in range(1,h-1):
        for j in range(1,w-1):
            if(bin_img[i][j]==0):
                contour_img[i][j] = 0
                sum = 0
                sum += bin_img[i - 1][j + 1]
                sum += bin_img[i][j + 1]
                sum += bin_img[i + 1][j + 1]
                sum += bin_img[i - 1][j]
                sum += bin_img[i + 1][j]
                sum += bin_img[i - 1][j - 1]
                sum += bin_img[i][j - 1]
                sum += bin_img[i + 1][j - 1]
                if sum ==  0:
                    contour_img[i][j] = 255

    return contour_img

if __name__ == "__main__":
    img_name = "./20210808/sample3.png"
    img = cv2.imread(img_name)
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    bin_img = get_binary_img(gray_img)
    contour_img = get_contour(bin_img)

    cv2.imshow("src", img)
    cv2.imshow("contour",contour_img)
    cv2.waitKey(0)

max_stars_repo_name: sgzqc/wechat

max_issues_repo_name: sgzqc/wechat

max_forks_repo_name: sgzqc/wechat


text:
# -*- coding: utf-8 -*-
"""
Codes by @kentaroy47

"""
try:
    # for tensorflow2
    #import tensorflow.keras.backend as K
    # Doesn't work ;(
    # https://stackoverflow.com/questions/61056781/typeerror-tensor-is-unhashable-instead-use-tensor-ref-as-the-key-in-keras
    import tensorflow.compat.v1.keras.backend as K

except:
    # tensorflow1 and keras 2.x
    import keras.backend as K
import numpy as np

# bunch of cal per layer
def count_linear(layers):
    MAC = layers.output_shape[1] * layers.input_shape[1]
    if layers.get_config()["use_bias"]:
        ADD = layers.output_shape[1]
    else:
        ADD = 0
    return MAC*2 + ADD

def count_conv2d(layers, log = False):
    if log:
        print(layers.get_config())
    # number of conv operations = input_h * input_w / stride = output^2
    numshifts = int(layers.output_shape[1] * layers.output_shape[2])

    # MAC/convfilter = kernelsize^2 * InputChannels * OutputChannels
    MACperConv = layers.get_config()["kernel_size"][0] * layers.get_config()["kernel_size"][1] * layers.input_shape[3] * layers.output_shape[3]

    if layers.get_config()["use_bias"]:
        ADD = layers.output_shape[3]
    else:
        ADD = 0

    return MACperConv * numshifts * 2 + ADD

def profile(model, log=False):
    # make lists
    layer_name = []
    layer_flops = []
    # TODO: relus
    inshape = []
    weights = []
    # run through models
    for layer in model.layers:
        if "dense" in layer.get_config()["name"] or "fc" in layer.get_config()["name"]:
            layer_flops.append(count_linear(layer))
            layer_name.append(layer.get_config()["name"])
            inshape.append(layer.input_shape)

            # weights seems to be broken in tf2.x
            try:
                weights.append(int(np.sum([K.count_params(p) for p in set(layer.trainable_weights)])))
            except:
                pass
        elif "conv" in layer.get_config()["name"] and "pad" not in layer.get_config()["name"] and "bn" not in layer.get_config()["name"] and "relu" not in layer.get_config()["name"] and "concat" not in layer.get_config()["name"]:
            layer_flops.append(count_conv2d(layer,log))
            layer_name.append(layer.get_config()["name"])
            inshape.append(layer.input_shape)
            try:
                weights.append(int(np.sum([K.count_params(p) for p in set(layer.trainable_weights)])))
            except:
                pass
        elif "res" in layer.get_config()["name"] and "branch" in layer.get_config()["name"]:
            layer_flops.append(count_conv2d(layer,log))
            layer_name.append(layer.get_config()["name"])
            inshape.append(layer.input_shape)
            try:
                weights.append(int(np.sum([K.count_params(p) for p in set(layer.trainable_weights)])))
            except:
                pass

    return layer_name, layer_flops, inshape, weights


text:
import numpy as np
import pickle
class NLinearModels(object):
    def __init__(self,x_example,number_of_regressors=4,learning_rate = 0.1):
        shape_input = x_example.reshape(-1).shape[0]
        limit = np.sqrt(6 / (shape_input + number_of_regressors))
        self.W = np.random.uniform(-limit,limit, size=(shape_input,number_of_regressors)) #HE INITIALIZATION
        #self.W = np.ones((shape_input,number_of_regressors))/10 #HE INITIALIZATION
        self.bias = np.zeros(number_of_regressors)
        self.learning_rate = learning_rate

    def forward(self,x):
        return x.dot(self.W) + self.bias

    def predict(self,x):
        x = np.array(x)
        x = x.reshape(x.shape[0],-1)
        return self.forward(x)

    def cost(self,y_hat,y):
        return ((y_hat-y)**2).mean(axis=0)

    def backward(self,x,y_hat,y):
        m = y_hat.shape[0]
        dl = 2*(y_hat-y)/m
        self.bias_gradient = np.sum(dl,axis=0)
        self.W_gradient = x.T.dot(dl)/m

    def train_on_batch(self,_input,target):
        _input = np.array(_input)
        y = np.array(target)
        x = _input.reshape(_input.shape[0],-1)
        y_hat = self.forward(x)
        cost = self.cost(y_hat,y).sum()
        self.backward(x,y_hat,y)
        self.update_weights()
        return cost

    def update_weights(self):

        self.W -= self.learning_rate * self.W_gradient
        self.bias -= self.learning_rate * self.bias_gradient

    def load(self):
        W = np.load(open('save_rl/W.npy',"rb"))
        bias = np.load(open('save_rl/bias.npy',"rb"))
        self.W = W
        self.bias = bias

    def save(self):
        np.save(open("save_rl/W.npy","wb"),self.W)
        np.save(open("save_rl/bias.npy","wb"),self.bias)

class ExperienceReplay(object):
    """
    During gameplay all the experiences < s, a, r, s’ > are stored in a replay memory.
    In training, batches of randomly drawn experiences are used to generate the input and target for training.
    """
    def __init__(self, max_memory=100, discount=.9):
        """
        Setup
        max_memory: the maximum number of experiences we want to store
        memory: a list of experiences
        discount: the discount factor for future experience

        In the memory the information whether the game ended at the state is stored seperately in a nested array
        [...
        [experience, game_over]
        [experience, game_over]
        ...]
        """
        self.max_memory = max_memory
        self.memory = list()
        self.discount = discount

    def remember(self, states, game_over):
        #Save a state to memory
        self.memory.append([states, game_over])
        #We don't want to store infinite memories, so if we have too many, we just delete the oldest one
        if len(self.memory) > self.max_memory:
            del self.memory[0]

    def get_batch(self, model, batch_size=10):

        #How many experiences do we have?
        len_memory = len(self.memory)

        #Calculate the number of actions that can possibly be taken in the game
        num_actions = 4

        #Dimensions of the game field
        env_dim = list(self.memory[0][0][0].shape)
        env_dim[0] = min(len_memory, batch_size)

        #We want to return an input and target vector with inputs from an observed state...
        inputs = np.zeros(env_dim)
        #...and the target r + gamma * max Q(s’,a’)
        #Note that our target is a matrix, with possible fields not only for the action taken but also
        #for the other possible actions. The actions not take the same value as the prediction to not affect them
        Q = np.zeros((inputs.shape[0], num_actions))

        #We draw states to learn from randomly
        for i, idx in enumerate(np.random.randint(0, len_memory,
                                                  size=inputs.shape[0])):
            """
            Here we load one transition <s, a, r, s’> from memory
            state_t: initial state s
            action_t: action taken a
            reward_t: reward earned r
            state_tp1: the state that followed s’
            """
#            idx = -1
            state, action_t, reward_t, state_tp1 = self.memory[idx][0]
            #We also need to know whether the game ended at this state
            game_over = self.memory[idx][1]

            #add the state s to the input
            inputs[i:i+1] = state
            # First we fill the target values with the predictions of the model.
            # They will not be affected by training (since the training loss for them is 0)
            Q[i] = model.predict([state])[0]

            """
            If the game ended, the expected reward Q(s,a) should be the final reward r.
            Otherwise the target value is r + gamma * max Q(s’,a’)
            """
            #if the game ended, the reward is the final reward
            if game_over:  # if game_over is True
                Q[i, action_t] = reward_t
            else:
                # r + gamma * max Q(s’,a’)
                next_round = model.predict([state_tp1])[0]
                Q[i, action_t] = reward_t + self.discount*np.max(next_round)
        return inputs, Q

    def load(self):
        self.memory = pickle.load(open("save_rl/memory.pkl","rb"))
    def save(self):
        pickle.dump(self.memory,open("save_rl/memory.pkl","wb"))

max_stars_repo_name: isival/IntroToAI

max_issues_repo_name: isival/IntroToAI

max_forks_repo_name: isival/IntroToAI