python source code of batch

Jacinle-master
- jaclearn
  - datasets
    - image_classification
      - mnist.py
      - svhn.py
      - __init__.py
      - cifar.py
    - __init__.py
  - nlp
    - graph
      - dependency_visualizer
        render.py
        templates.py
        __init__.py
        utils.py
      - __init__.py
    - tree
      - ptb.py
      - constituency.py
      - traversal.py
      - node.py
      - __init__.py
    - __init__.py
    - sng_parser
      - __init__.py
  - logic
    - __init__.py
    - propositional
      - logic_induction
        logic_induction.pxd
        travis.sh
        setup.py
        logic_induction.pyx
        __init__.py
        csrc
        logic_induction_impl.cc
        logic_induction.h
        .gitignore
      - __init__.py
  - models
    - naive_bayes
      - hybrid_nb.py
      - __init__.py
    - __init__.py
  - rl
    - env.py
    - simulator
      - controller.py
      - __init__.py
      - pack.py
    - engines
      - mujoco
        gym_recorder.py
        __init__.py
      - __init__.py
    - algo
      - advantage.py
      - __init__.py
      - math.py
    - space.py
    - __init__.py
    - proxy.py
    - envs
      - simple
        move_right.py
        __init__.py
      - gym.py
      - gym_adapter.py
      - __init__.py
      - maze
        taxi.py
        _maze_visualizer.py
        __init__.py
        maze.py
      - nintendo
        _wrapper.py
        mario.py
        __init__.py
  - dataflow
    - collections.py
    - batch.py
    - __init__.py
    - utils.py
    - dataflow.py
  - imageaug
    - photography.py
    - executor.py
    - shape.py
    - __init__.py
    - cblk.py
  - visualize
    - imgrid.py
    - box.py
    - html_table.py
    - plot.py
    - __init__.py
    - imshow.py
  - data
    - sampler.py
    - stat.py
    - __init__.py
  - mldash
    - __init__.py
  - __init__.py
  - math
    - rotation.py
    - __init__.py
  - vision
    - graph_cut.py
    - graph_cut_inpaint.py
    - __init__.py
    - coco
      - src
        maskApi.h
        maskApi.c
        .gitignore
        gason.cpp
        gason.h
      - travis.sh
      - mask_utils.py
      - setup.py
      - pycocotools
        coco.py
        mask.py
        cocoeval.py
        __init__.py
        _mask.pyx
      - __init__.py
      - .gitignore
    - patch_match.py
  - embedding
    - visualize_tb.py
    - word_embedding.py
    - embedding_utils.py
    - __init__.py
    - constant.py
- jactf
  - utils
    - init.py
    - __init__.py
  - __init__.py
- examples
  - torch-starter
    - trainval.py
    - test.py
    - extract-coco-features.py
  - comm-broadcast-gather
    - broadcast.py
    - gather.py
  - cython
    - csrc
      - loop.c
    - jit_example.py
  - utils-ipdb
    - main.py
  - comm-cs
    - client.py
    - server.py
  - web-starter
    - main.py
    - app
      - index.py
      - templates
        index.html
      - __init__.py
  - utils-defaults
    - default_args.py
  - torch-dataloader
    - multi_gpu_wrapper.py
    - syncable_dataset.py
  - comm-distrib
    - zmq_benchmark.py
    - out.py
    - in_benchmark.py
    - in.py
  - learn-logic
    - propositional-induction.py
  - python-starter
    - script.py
  - configs
    - test-configs-envrion2.py
  - dependency-visualizer
    - example.py
    - .gitignore
  - comm-service
    - client.py
    - server.py
  - torch-jacvision
    - test-jv.py
- jacnp
  - indexing.py
  - shape.py
  - batch.py
  - nd.py
  - __init__.py
- .gitmodules
- jacweb
  - web
    - application.py
    - __init__.py
  - __init__.py
  - session
    - memcached.py
    - __init__.py
    - session.py
- LICENSE
- vendors
  - PyPatchMatch
  - AdvancedIndexing-PyTorch
  - SceneGraphParser
  - PreciseRoIPooling
  - Synchronized-BatchNorm-PyTorch
  - pytorch_reservoir
    - torch_sampling
    - pytorch_reservoir.py
    - travis.sh
  - PyGCO
    - pygco
    - pygco_inpaint
  - JacMLDash
- jacinle.yml
- .vim-template:.py
- jactorch
  - graph
    - variable.py
    - __init__.py
    - parameter.py
  - utils
    - init.py
    - grad.py
    - unittest.py
    - meta.py
    - __init__.py
  - functional
    - quantization.py
    - grad.py
    - probability.py
    - sample.py
    - loglinear.py
    - kernel.py
    - indexing.py
    - shape.py
    - mask.py
    - linalg.py
    - __init__.py
    - arith.py
    - meshgrid.py
  - optim
    - optimizer_group.py
    - __init__.py
    - custom_optimizer.py
    - accum_grad.py
    - quickaccess.py
    - weight_decay.py
  - cli.py
  - cuda
    - copy.py
    - __init__.py
  - models
    - __init__.py
    - vision
      - vgg.py
      - resnet.py
      - __init__.py
  - nn
    - container.py
    - simple.py
    - probability.py
    - embedding.py
    - residual.py
    - normalization.py
    - rnn_layers.py
    - sync_batchnorm
      - __init__.py
    - losses
      - losses.py
      - functional.py
      - __init__.py
    - functional.py
    - prroi_pool
      - functional.py
      - __init__.py
      - .gitignore
    - cnn
      - conv.py
      - functional.py
      - layers.py
      - coord_conv.py
      - __init__.py
    - mask
      - functional.py
      - __init__.py
    - word_embedding.py
    - rnn_utils.py
    - __init__.py
    - neural_logic
      - recurrent_layer.py
      - layer.py
      - __init__.py
      - modules
        dimension.py
        neural_logic.py
        _utils.py
        __init__.py
    - linear.py
    - quickaccess.py
    - gumbel_softmax
      - softmax.py
      - functional.py
      - __init__.py
  - quickstart
    - models.py
    - train.py
    - __init__.py
    - inference.py
  - data
    - layout.py
    - collate
      - collate_v3.py
      - collate_v2.py
      - __init__.py
      - utils.py
      - collate_v1.py
    - dataloader
      - dataloader.py
      - __init__.py
      - dataloader_torch030.py
    - __init__.py
    - dataset.py
  - parallel
    - dict_gather.py
    - comm.py
    - replicate.py
    - data_parallel.py
    - replication_callback.py
    - __init__.py
    - user_scattered.py
  - io.py
  - __init__.py
  - serve
    - service.py
  - vision
    - conv.py
    - morphology.py
    - geometry.py
    - _utils.py
    - peak.py
    - gradient.py
    - __init__.py
    - smooth.py
  - train
    - env.py
    - monitor.py
    - __init__.py
    - utils.py
    - tb.py
  - transforms
    - coor
      - functional.py
      - __init__.py
      - transforms.py
    - bbox
      - functional.py
      - __init__.py
      - transforms.py
    - __init__.py
    - vision
      - functional
        image.py
        coor.py
        _utils.py
        bbox.py
        __init__.py
      - __init__.py
      - transforms.py
    - image
      - functional.py
      - __init__.py
      - transforms.py
- README.md
- tests
  - test-utils-debug.py
  - test-torch-trainer.py
  - test-learn-ptb.py
  - test-torch-conv.py
  - test-torch-inferencer.py
  - test-torch-dataloader.py
- scripts
  - docs-server.py
  - name-server.py
  - inspect-file.py
  - format-header.py
- bin
  - jac-run
  - jac-script.sh
  - jac-mldash
  - jac-inspect-file
  - _jac-init.sh
  - jac-update
  - jac-docs-server
  - jac-format-header
  - jac-crun
  - _jac-init-gen.py
  - jac-name-server
- requirements.txt
- .gitignore
- docs
  - Makefile
  - source
    - jaclearn.nlp.graph.rst
    - jactorch.nn.mask.rst
    - jaclearn.datasets.image_classification.rst
    - jactorch.transforms.vision.functional.rst
    - jacinle.jit.cext.rst
    - jacinle.rst
    - jaclearn.datasets.rst
    - jaclearn.vision.coco.pycocotools.rst
    - jacinle.logging.rst
    - jactorch.nn.gumbel_softmax.rst
    - jaclearn.vision.coco.rst
    - jacinle.random.rst
    - jaclearn.rl.envs.nintendo.rst
    - jactorch.models.rst
    - jaclearn.embedding.rst
    - jactorch.data.rst
    - jacinle.concurrency.rst
    - jactorch.transforms.rst
    - jaclearn.mldash.rst
    - jacinle.jit.cython.rst
    - jactorch.parallel.rst
    - jacinle.comm.rst
    - jactorch.data.collate.rst
    - jaclearn.dataflow.rst
    - jaclearn.data.rst
    - jaclearn.rl.algo.rst
    - jacinle.cli.rst
    - jactorch.nn.losses.rst
    - jacinle.utils.rst
    - jactorch.utils.rst
    - jactorch.graph.rst
    - jaclearn.nlp.graph.dependency_visualizer.rst
    - jactorch.nn.neural_logic.rst
    - jacinle.storage.kv.rst
    - jacinle.comm.distrib.rst
    - jacweb.session.rst
    - jactorch.quickstart.rst
    - jaclearn.models.naive_bayes.rst
    - jactorch.models.vision.rst
    - jaclearn.visualize.rst
    - jaclearn.rst
    - jactorch.nn.neural_logic.modules.rst
    - jactf.rst
    - jaclearn.rl.envs.simple.rst
    - jaclearn.nlp.tree.rst
    - jactorch.transforms.image.rst
    - jacnp.rst
    - jacinle.image.rst
    - jaclearn.rl.envs.rst
    - jaclearn.nlp.rst
    - jaclearn.imageaug.rst
    - jactorch.nn.sync_batchnorm.rst
    - jacinle.config.rst
    - jacweb.rst
    - jacinle.event.rst
    - jaclearn.rl.rst
    - jactorch.nn.rst
    - jactorch.data.dataloader.rst
    - jaclearn.vision.rst
    - modules.rst
    - jactorch.nn.cnn.rst
    - jaclearn.nlp.sng_parser.rst
    - jactorch.functional.rst
    - jactf.utils.rst
    - index.rst
    - jacinle.io.rst
    - conf.py
    - jactorch.transforms.coor.rst
    - jactorch.nn.prroi_pool.rst
    - jacinle.jit.rst
    - jactorch.transforms.vision.rst
    - jaclearn.rl.envs.maze.rst
    - jactorch.vision.rst
    - jacinle.storage.rst
    - jacweb.web.rst
    - jaclearn.rl.simulator.rst
    - jactorch.train.rst
    - jactorch.transforms.bbox.rst
    - jaclearn.models.rst
    - jactorch.optim.rst
    - jactorch.rst
    - jactorch.cuda.rst
  - requirements.txt
  - .gitignore
- jacinle
  - utils
    - init.py
    - container.py
    - imp.py
    - cache.py
    - env.py
    - naming.py
    - inspect.py
    - argument.py
    - printing.py
    - matching.py
    - vendor.py
    - registry.py
    - meta.py
    - _inspect_mate.py
    - __init__.py
    - numeric.py
    - filelock.py
    - exception.py
    - deprecated.py
    - defaults.py
    - meter.py
    - context.py
    - value_scheduler.py
    - debug.py
    - network.py
    - tqdm.py
    - enum.py
  - storage
    - kv
      - kv.py
      - memcached.py
      - __init__.py
      - mem.py
      - lmdb.py
    - __init__.py
  - comm
    - broadcast.py
    - gather.py
    - cs.py
    - distrib
      - name_server.py
      - pipe.py
      - _configs.py
      - controller.py
      - __init__.py
    - __init__.py
    - echo.py
    - service.py
  - jit
    - cython
      - __init__.py
      - jit_compile.py
    - __init__.py
    - cext
      - travis.py
      - __init__.py
  - cli
    - git.py
    - keyboard.py
    - argument.py
    - __init__.py
    - device.py
  - concurrency
    - pool.py
    - future.py
    - counter.py
    - zmq_utils.py
    - queue.py
    - event.py
    - __init__.py
    - process.py
    - shmarray.py
    - packing.py
  - config
    - environ.py
    - environ_v2.py
    - g.py
    - __init__.py
  - event
    - registry.py
    - __init__.py
  - io
    - pretty.py
    - common.py
    - __init__.py
    - fs.py
    - network.py
    - tempfile.py
  - __init__.py
  - random
    - _lazy_bind.py
    - __init__.py
    - rng.py
  - image
    - codecs.py
    - imgio.py
    - __init__.py
    - backend.py
    - imgproc.py
  - logging
    - logger.py
    - __init__.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : batch.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/18/2018
#
# This file is part of Jacinle.
# Distributed under terms of the MIT license.

import collections
import numpy as np

__all__ = ['batchify', 'unbatchify']


def batchify(inputs):
    first = inputs[0]
    if isinstance(first, (tuple, list, collections.UserList)):
        return [batchify([ele[i] for ele in inputs]) for i in range(len(first))]
    elif isinstance(first, (collections.Mapping, collections.UserDict)):
        return {k: batchify([ele[k] for ele in inputs]) for k in first}
    return np.stack(inputs)


def unbatchify(inputs):
    if isinstance(inputs, (tuple, list, collections.UserList)):
        outputs = [unbatchify(e) for e in inputs]
        return list(map(list, zip(*outputs)))
    elif isinstance(inputs, (collections.Mapping, collections.UserDict)):
        outputs = {k: unbatchify(v) for k, v in inputs.items()}
        first = outputs[0]
        return [{k: outputs[k][i] for k in inputs} for i in range(len(first))]
    return list(inputs)